code: def test_shared_modules_pure(self):
      class Block(nnx.Module):
          def __init__(self, linear: nnx.Linear, *, rngs: nnx.Rngs):
              self.linear = linear
              self.bn = nnx.BatchNorm(2, rngs=rngs)
  
          def __call__(self, x):
              x = self.linear(x)
              x = self.bn(x)
              return nnx.relu(x)
  
      class Model(nnx.Module):
          def __init__(self, *, rngs: nnx.Rngs):
              shared = nnx.Linear(2, 2, rngs=rngs)
              self.block1 = Block(shared, rngs=rngs)
              self.block2 = Block(shared, rngs=rngs)
  
          def __call__(self, x):
              x = self.block1(x)
              x = self.block2(x)
              return x
  
      @jax.jit
      def train_step(
          state: nnx.State, graphdef: nnx.GraphDef[Model], x, y
      ):
          model = graphdef.merge(state)
  
          @nnx.grad
          def loss_fn(model: Model):
              with nnx.flags(use_running_average=False):
                  y_pred = model(x)
              return jnp.mean((y - y_pred) ** 2)
  
          grads = loss_fn(model)
          model.update(
              jax.tree_map(
                  lambda w, g: w - 0.1 * g,
                  model.extract(nnx.Param),
                  grads,
              )
          )
  
          return model.split()
  
      graphdef: nnx.GraphDef[Model]
      state, graphdef = Model(rngs=nnx.Rngs(0)).split()
  
      x = np.random.uniform(size=(4, 2))
      y = np.random.uniform(size=(4, 2))
  
      for _i in range(3):
          state, graphdef = train_step(state, graphdef, x, y)
  
      model = graphdef.merge(state)
  
      assert model.block1.linear.bias is not None
      assert model.block2.linear.bias is not None
      assert model.block1.linear.kernel is model.block2.linear.kernel
      assert model.block1.linear.bias is model.block2.linear.bias
      assert model.block1.bn is not model.block2.bn
  
file_path: flax/flax/experimental/nnx/tests/test_integration.py
imports: import typing
  import jax
  import jax.numpy
  import numpy
  from flax.experimental import nnx
label_desc: Use image node_id 2 for calling the TestIntegration obj's underlying member method code with example usage: obj.test_shared_modules_pure() and returns: model, nnx, x, jnp