How to clone a state in flax.nnx?

[JINKEHE]

Hello,

I have two questions about how to clone things in flax.nnx.

1. How to clone a state?

   • I see that we can clone an nnx model with cloned_model = nnx.clone(model). But is there anything like graphdef, state = nnx.split(model). new_state = nnx.clone(state)?

2. I also have a question about how clone is implemented in flax.nnx.
   in graph.py, we have

def clone(node: Node) -> Node:
  """Create a deep copy of the given graph node.

  Example usage::

    >>> from flax import nnx

    >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
    >>> cloned_model = nnx.clone(model)
    >>> model.bias.value += 1
    >>> assert (model.bias.value != cloned_model.bias.value).all()

  Args:
    node: A graph node object.
  Returns:
    A deep copy of the :class:`Module` object.
  """
  graphdef, state = split(node)
  return merge(graphdef, state)

But why does this implement a deepcopy? Is it in split or merge? Does it mean that everytime we use the nnx functional API, we perform a deepcopy of neural network weights? This might cause some concerns on performance ...

Thanks in advance for your help!

[cgarciae]

Hey @JINKEHE

1. Since State is a Pytree you can clone it using jax.tree.map:

state = jax.tree.map(lambda x: x, state) # clone

2. split and merge don't use python's deepcopy but do define a complete traversal of the object graph for all NNX objects and JAX pytrees. Think of split as jax.tree.flatten and merge as jax.tree.unflatten. In fact, NNX also has nnx.graph.flatten and nnx.graph.unflatten which are used by split and merge under the hood.

[JINKEHE]

Thank you!