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!

[DBraun]

with Flax 0.11.1 and jaxlib 0.7.0, I have a related question about storing self.layers = [nnx.Linear(2, 2, rngs=rngs)]. It seems to not work well with checkpointing.

import tempfile
from flax import nnx
import orbax.checkpoint as ocp


class ModelWithLists(nnx.Module):
    def __init__(self, i):
        self.layers = [
            nnx.Linear(2, 2, rngs=nnx.Rngs(i+j))
            for j in range(3)
        ]
        self.linear = nnx.Linear(4, 4, rngs=nnx.Rngs(778))


model = ModelWithLists(0)
print(model.layers[0].kernel.value)  # 1. Original weights

with tempfile.TemporaryDirectory() as tmpdir:
    # Save checkpoint
    state = nnx.state(model)
    mngr = ocp.CheckpointManager(tmpdir, checkpointers={"state": ocp.StandardCheckpointer()})
    mngr.save(1, items={"state": state})
    print("Checkpoint saved")

    model2 = ModelWithLists(1)  # Create different weights
    print(model2.layers[0].kernel.value) # 2. Should be different from 1.

    # Restore checkpoint
    restored = mngr.restore(1)
    saved_state = restored['state']

    nnx.update(model2, saved_state)
    print(model2.layers[0].kernel.value)
    # print(model2.layers.value["value"][0]["kernel"]["value"])  # 3. Should be same as 1.

Output:

[[ 0.75672996 -0.6857006 ]
 [-0.5688339  -0.8757607 ]]
WARNING:absl:Configured `CheckpointManager` using deprecated legacy API. Please follow the instructions at https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html to migrate.
Checkpoint saved
[[ 0.40635777 -0.55007243]
 [ 0.68637383 -0.71790355]]
WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.
Traceback (most recent call last):
  File "/mnt/c/Users/admin/GitHub/flax_question/flax_checkpoint_minimal.py", line 37, in <module>
    nnx.update(model2, saved_state)
  File "/home/admin/.local/lib/python3.12/site-packages/flax/nnx/graph.py", line 2406, in update
    _graph_update_dynamic(node, state)
  File "/home/admin/.local/lib/python3.12/site-packages/flax/nnx/graph.py", line 1464, in _graph_update_dynamic
    _graph_update_dynamic(current_value, value)
  File "/home/admin/.local/lib/python3.12/site-packages/flax/nnx/graph.py", line 1447, in _graph_update_dynamic
    raise ValueError(
ValueError: Cannot set key '0' on immutable node of type list

I can get the nnx.update to work if self.layers is instead

        self.layers = nnx.Param([
            nnx.Linear(2, 2, rngs=nnx.Rngs(i+j))
            for j in range(3)
        ])

But this is also kind of ugly, and then print(model2.layers[0].kernel.value) no longer works.
So I have to do print(model2.layers.value["value"][0]["kernel"]["value"])
Final output:

[[ 0.75672996 -0.6857006 ]
 [-0.5688339  -0.8757607 ]]
WARNING:absl:Configured `CheckpointManager` using deprecated legacy API. Please follow the instructions at https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html to migrate.
Checkpoint saved
[[ 0.40635777 -0.55007243]
 [ 0.68637383 -0.71790355]]
WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.
[[ 0.75672996 -0.6857006 ]
 [-0.5688339  -0.8757607 ]]

[DBraun]

Avoiding the deprecated legacy API fixed the issue above. Here's a working example:

import tempfile
from flax import nnx
import orbax.checkpoint as ocp


class ModelWithLists(nnx.Module):
    def __init__(self, i):
        self.layers = [
            nnx.Linear(2, 2, rngs=nnx.Rngs(i+j))
            for j in range(3)
        ]
        self.linear = nnx.Linear(4, 4, rngs=nnx.Rngs(778))


model = ModelWithLists(0)
print(model.layers[0].kernel.value)  # 1. Original weights

with tempfile.TemporaryDirectory() as tmpdir:
    # Save checkpoint
    mngr = ocp.CheckpointManager(tmpdir, options=ocp.CheckpointManagerOptions(enable_async_checkpointing=False))
    save_args = ocp.args.Composite(
        model=ocp.args.PyTreeSave(nnx.state(model)),
    )
    mngr.save(1, args=save_args)
    print("Checkpoint saved")

    model2 = ModelWithLists(1)  # Create different weights
    print(model2.layers[0].kernel.value) # 2. Should be different from 1.

    # Restore checkpoint
    restore_args = ocp.args.Composite(
        model=ocp.args.PyTreeRestore(nnx.state(model)),
    )
    restored = mngr.restore(1, args=restore_args)

    nnx.update(model2, restored.model)
    print(model2.layers[0].kernel.value)

[rademacher-p]

@cgarciae Can I have some further guidance on this, given the changes to Flax in the past year?

As per our discussion above, I've been wrapping any arbitrary/custom pytrees in Variable and using the value property for get/set. Now, value is deprecated and a refactor to get_value/set_value would be a bit syntactically unpleasant in places. If I use the new nnx.data instead of Variable, the attribute get/set syntax is nice, but I get the same nnx.update error as above.

Are there any other ways to proceed besides get_value/set_value? Changing all my pytrees into Module would seem heavyweight.

One comparable case I've seen in the source is for nnx.Optimizer, where the optax state pytree is transformed via to_opt_state and nnx.data; updates to the resultant opt_state attribute are performed with nnx.update. Is there a good reason to do something like this, rather than just the Variable wrap I'm frequently performing?

Thanks for the help 🙏

[cgarciae]

but yeah try to wrap leaf values in Variables, do your operations normally on the pytree, use nnx.update to update previous Variable references.

[cgarciae]

alternatively you can also just do the updates in-place if you have control of this e.g:

a = A()

def update_fn(v):
    v[...] += 1

jax.tree.map(update_fn, a, is_leaf=lambda x: isinstance(x, nnx.Variable))

print(a)

[rademacher-p]

@cgarciae That will work, but may not be generally applicable - specifically, if I want to use Pytrees from some third-party library.

Is there a downside to self.foo = nnx.Variable(Node())? Should I follow the example of nnx.Optimizer and map over the Node leaves, transforming them to Variable?

So it looks like your recommendation is to follow the example of nnx.Optimizer. To be clear, do you think there is a problem with doing nnx.Variable(Node())?

[cgarciae]

do you think there is a problem with doing nnx.Variable(Node())

@rademacher-p originally Optimizer was written like this as well e.g. self.opt_state = OptState(opt_state). We ran into an edge case eventually (don't exactly remember) so we switched to wrapping the leaves instead but you can try. It does work:

@jax.tree_util.register_dataclass
@dataclasses.dataclass
class Node:
    var: tp.Any

class A(nnx.Pytree):
    def __init__(self):
        self.foo = nnx.Variable(Node(jnp.int32(0)))

a = A()
new_a = jax.tree.map(lambda x: x + 1, a)
nnx.update(a, nnx.state(new_a))

print(a)

[rademacher-p]

OK, thanks for the info. I'll probably just keep my non-Array Variables as is until value is removed or I run into that edge case myself 😅

[cgarciae]

Oh yeah the issue was that we wanted to propagate Param metadata to OptState metadata, but if you don't need this wrapping the entire structure should work.

[cgarciae]

@rademacher-p do you need .value for only for access or for setting it too? We might be able to keep .value for access only long term if its useful.

[rademacher-p]

Lots of attribute setting after sending these custom trees through pure JAX funcs.

I'll probably just do the refactor w/ set_value for most of my variables. Others, such as a Variable containing State of target RL policy weights, I might try the Optimizer approach and map the leaves.