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)