code: def test_grad_example(self):
      m = nnx.Dict(
          a=nnx.Sequence([nnx.Param(1.0), nnx.BatchStat(-10)]),
          b=nnx.Param(2.0),
          c=100,
      )
  
      params = m.extract(nnx.Param)
  
      def loss(params):
          return sum(2 * p for p in jax.tree_util.tree_leaves(params))
  
      grads = jax.grad(loss)(params)
      m.update(grads)
  
      assert m.a[0] == 2.0
      assert m.a[1] == -10
      assert m.b == 2.0
      assert m.c == 100
  
file_path: flax/flax/experimental/nnx/tests/test_partitioning.py
imports: import jax
  import pytest
  from flax.experimental import nnx
label_desc: Use image node_id 8 for calling the TestPartitioning obj's underlying member method code with example usage: obj.test_grad_example() and returns: sum