Is there a way to give pytree state to a NNX module?

[wittenator]

I want to include the dense solution of a diffrax ODE in a NNX module. Since this is a pytree object with arrays, nnx.jit complains about array root leafs. Is there a way to tell nnx to simply ignore this?

[cgarciae]

Hi @wittenator, Array leaves are already supported as of #4612, please update flax to the latest version.

[wittenator]

Oh that's awesome! Are there any limitations to this? Using the diffrax object seems to work when working without nnx.jit, but there seems to be a problem, when I activate it. Do you know if this is an nnx problem?

Traceback (most recent call last):
  File "/repo/.venv/lib/python3.12/site-packages/hydra_zen/wrapper/_implementations.py", line 439, in __call__
    return func(
           ^^^^^
  File "/repo/scripts/train.py", line 156, in main
    train(
  File "/repo/scripts/train.py", line 106, in train
    cached_train_step(
  File "/repo/.venv/lib/python3.12/site-packages/flax/nnx/graph.py", line 1311, in cache_args_wrapper
    return f(*cached_args, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/repo/.venv/lib/python3.12/site-packages/flax/nnx/transforms/compilation.py", line 341, in jit_wrapper
    pure_args, pure_kwargs = extract.to_tree(
                             ^^^^^^^^^^^^^^^^
  File "/repo/.venv/lib/python3.12/site-packages/flax/nnx/extract.py", line 280, in to_tree
    return jax.tree.map(
           ^^^^^^^^^^^^^
  File "/repo/.venv/lib/python3.12/site-packages/jax/_src/tree.py", line 155, in map
    return tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/repo/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 358, in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/repo/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 358, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
                             ^^^^^^
  File "/repo/.venv/lib/python3.12/site-packages/flax/nnx/extract.py", line 281, in <lambda>
    lambda x: split_fn(split_ctx, (), prefix, x)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/repo/.venv/lib/python3.12/site-packages/flax/nnx/transforms/compilation.py", line 101, in _jit_split_fn
    return extract.NodeStates.from_split(*ctx.flatten(x, with_paths=False))
                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/repo/.venv/lib/python3.12/site-packages/flax/nnx/graph.py", line 1466, in flatten
    variable.raw_value for variable in node_static_cache.variables
    ^^^^^^^^^^^^^^^^^^
AttributeError: 'jaxlib.xla_extension.ArrayImpl' object has no attribute 'raw_value'. Did you mean: '_value'?

[wittenator]

Ah hmm, seems like it is connected to using nnx.cached_partial

[cgarciae]

Yeah sorry, cached_partial doesn't yet support Array leaves.