code: def test_train_state_methods(self):
      class Foo(nnx.Module):
          def __init__(self, *, rngs: nnx.Rngs):
              self.linear = nnx.Linear(2, 4, rngs=rngs)
              self.batch_norm = nnx.BatchNorm(4, rngs=rngs)
  
          def __call__(self, x: jax.Array, train: bool) -> jax.Array:
              x = self.linear(x)
              x = self.batch_norm(x, use_running_average=not train)
              return x
  
      module = Foo(rngs=nnx.Rngs(0))
      params, batch_stats, graphdef = module.split(
          nnx.Param, nnx.BatchStat
      )
  
      state = nnx.TrainState(
          graphdef,
          params=params,
          tx=optax.sgd(1.0),
          batch_stats=batch_stats,
      )
  
      x = jax.numpy.ones((1, 2))
      y, _updates = state.apply("params", "batch_stats")(x, train=True)
  
      assert y.shape == (1, 4)
  
      # fake gradient
      grads = jax.tree_map(jnp.ones_like, state.params)
      # test apply_gradients
      state = state.apply_gradients(grads)
  
imports: import jax
  import jax.numpy
  import optax
  from flax.experimental import nnx
label_desc: Use image node_id 2 for calling the TestHelpers obj's underlying member method code with example usage: obj.test_train_state_methods() and returns: x