Do JAX Arrays in a Module always have to be Variables?

[rademacher-p]

Is this possible? Not really a problem, but I have arrays that I know will remain static in my nnx.Module and would prefer to always exclude them from nnx.state, avoid unnecessary syntax (nnx.Variable(), var.value), etc.

[cgarciae]

Hi @rademacher-p, using jax.Arrays directly in your Modules is now supported, give it a try!

[rademacher-p]

@cgarciae OK, that helped me understand my issue better. Has something to do with my use of custom Pytree nodes. The snippet below raises an error:

import jax
import jax.numpy as jnp
from flax import nnx


@jax.tree_util.register_pytree_node_class
class Node:
    def __init__(self, var):
        self.var = var

    def tree_flatten(self):
        return (self.var,), ()

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(children[0])


class Model(nnx.Module):
    def __init__(self, foo):
        self.foo = foo


model = Model(Node(jnp.float32(0.0)))

nnx.update(model, nnx.state(model))

ValueError: Cannot set key 0 on immutable node of type Node

Can get around this by changing to self.foo = nnx.Variable(foo). Do you know if I can avoid the error without this kind of mod?

[cgarciae]

Right, its a bit hard to update arbitrary pytrees. Wrapping foo in Variable will work.

[rademacher-p]

OK, fair enough. Thanks for the help!