Are bool members of nnx.Module treated as static arguments when jitting the nnx.Module member function?

[HeavyCrab]

I know that if a function uses some global variable, jit won't throw an error, it would treat it as a side-effect instead. For example,

@nnx.jit
def f(x):
    if y:
        return x
    else:
        return x + 1
x = jnp.array(1)
y = False
print(f(x))
y = True
print(f(x))

This code will output

2
2

since the side effect will only be executed once.

However, I found things being different when jitting member functions of nnx.Module.
For example,

class mlp(nnx.Module):
    def __init__(self, in_dim, rngs):
        self.in_dim = in_dim
        self.layers = nnx.Linear(in_features=in_dim, out_features=in_dim, rngs=rngs)
        self.flag = True
        
    @nnx.jit
    def __call__(self, x):
        if self.flag:
            return x
        else:
            x = self.layers(x)
            return x

x = jnp.array([1, 2, 3])
rngs = nnx.Rngs(0)
model = mlp(3, rngs)

model.flag = True
print(forward(model, x))
model.flag = False
print(forward(model, x))

model.__call__ seems to behave correctly when relying on self.flag. The output is

[1 2 3]
[ 0.536822  -3.5292592  3.192309 ]

My question is:

Does that mean nnx.jit automatically treats all boolean members in nnx.Module as static arguments? Can I safely use if-else condition based on the boolean member in the nnx.Module class?

Thanks.

[cgarciae]

jit will only recompile if one of its inputs change, it doesn't track variables passed as a capture. To fix this have y be an input to f and mark y using static_argnums so JAX doesn't convert it to an array.

[HeavyCrab]

Thanks for the reply! I know that is the case for a global function. My question is: for a member function of nnx.Module, do we need to pass the attributes into the member function and manually mark it as static? In my case above, it seems the member function already tracks the object attribute self.flag as a static argument.

[cgarciae]

Hey! Sorry I might have stopped reading a bit early. I'll create a new response.

[cgarciae]

NNX will recognize NNX types (Module, Variable, etc) and pytree types (e.g. list, tuple, dict, etc) as nodes, and will treat any other type as static. If you change the value of static attributes jit will see this and recompile. This means you can safely use conditionals of static values.