[dparikh79]

What does this PR do?

Adds arrays=True to nnx.clone per #5461. When set, the underlying jax.Array buffers of each Variable are copied so the clone has independent memory and works with jit(donate_argnums=...). Default False so existing callers see no change. Setting arrays=True forces variables=True because a buffer copy without a fresh wrapper has nowhere to live.

Implementation: jax.tree.map(jax.numpy.copy, state) between split and merge. Used jnp.copy rather than jnp.array to preserve sharding on multi-device arrays.

Expanded the docstring to call out the default buffer-sharing per @samanklesaria's note that it should be better documented.

Two regression tests in tests/nnx/module_test.py. test_clone_arrays_distinct_buffers asserts unsafe_buffer_pointer() differs after arrays=True and matches for the default clone. test_clone_arrays_donate_argnums is @johnlyzhou's exact donate_argnums repro from the issue. Existing test_clone still passes. Full sweep: 77 passed / 1 skipped in module_test.py, 152 passed in graph_utils_test.py.

Fixes #5461

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).
• This change is discussed in a Github issue/discussion: nnx.clone creates buffer copies with the same IDs, causing errors with donate_argnums #5461
• The documentation and docstrings adhere to the documentation guidelines.
• This change includes necessary high-coverage tests.

[samanklesaria]

Looks good to me!

[cgarciae]

we should raise a ValueError if arrays=True and variables=False as this implementation currently does not support it.

[cgarciae]

since the jax.tree.map effectively creates copies of all the Variables then `you could actually set:

Suggested change

	variables = True
	variables = False

[vfdev-5]

@dparikh79 please update the PR according to the suggestions. Thanks!