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!