Copybara import of the project:

--
1f27faa by Cristian Garcia <cgarcia.e88@gmail.com>:

[nnx] improve nnx basics

PiperOrigin-RevId: 625789767

flax/experimental/nnx/__init__.py

@@ -33,15 +33,8 @@
 from .nnx.module import M as M
 from .nnx.module import Module as Module
 from .nnx.graph_utils import merge as merge
-from .nnx.graph_utils import full_merge as full_merge
 from .nnx.graph_utils import split as split
-from .nnx.graph_utils import full_split as full_split
 from .nnx.graph_utils import update as update
-from .nnx.graph_utils import full_update as full_update
-from .nnx.graph_utils import clone as clone
-from .nnx.graph_utils import pop as pop
-from .nnx.rnglib import init as init
-from .nnx.rnglib import empty as empty
 from .nnx.nn import initializers as initializers
 from .nnx.nn.activations import celu as celu
 from .nnx.nn.activations import elu as elu

flax/experimental/nnx/examples/toy_examples/09_parameter_surgery.py

@@ -52,10 +52,5 @@ def __call__(self, x):
 # split the parameters into trainable and non-trainable parameters
 trainable_params, non_trainable, static = model.split(is_trainable, ...)
 
-print(
-  'trainable_params =',
-  jax.tree_util.tree_map(jax.numpy.shape, trainable_params),
-)
-print(
-  'non_trainable = ', jax.tree_util.tree_map(jax.numpy.shape, non_trainable)
-)
+print('trainable_params =', jax.tree_util.tree_map(jax.numpy.shape, trainable_params))
+print('non_trainable = ', jax.tree_util.tree_map(jax.numpy.shape, non_trainable))