code: def test_train_state(self):
      m = nnx.Dict(a=nnx.Param(1), b=nnx.BatchStat(2))
  
      params, batch_stats, graphdef = m.split(nnx.Param, nnx.BatchStat)
  
      state = nnx.TrainState(
          graphdef,
          params=params,
          tx=optax.sgd(1.0),
          batch_stats=nnx.TreeNode(batch_stats),
          other=nnx.Variable(100),
          int=200,
      )
  
      leaves = jax.tree_util.tree_leaves(state)
  
      assert 1 in leaves
      assert 2 in leaves
      assert 100 in leaves
      assert 200 not in leaves
  
imports: import jax
  import jax.numpy
  import optax
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling the TestHelpers obj's underlying member method code with example usage: obj.test_train_state() without return types