code: def test_grad(self):
      p1 = nnx.Param(10.0)
      p2 = nnx.Param(20.0)
  
      m = nnx.Dict(
          a=nnx.Sequence([p1, p2]),
          b=p1,
          c=7,
          d=5.0,
      )
  
      @nnx.grad
      def f(m: nnx.Dict):
          # sum all params
          return m["a"][0] + m["a"][1] + m["b"]
  
      grads = f(m)
  
      assert m.a.variables["0"] is m.variables.b
      assert isinstance(grads, nnx.State)
      assert grads["a"]["0"] == 2.0
      assert isinstance(grads.a.variables["0"], nnx.Variable)
      assert grads["a"]["1"] == 1.0
      assert isinstance(grads.a.variables["1"], nnx.Variable)
      assert len(grads.flat_state()) == 2
  
      m.update(grads)
  
      assert m.a.variables["0"] is m.variables.b
      assert m["a"][0] == 2.0
      assert m["a"][1] == 1.0
      assert m["b"] == 2.0
      assert m["c"] == 7
      assert m["d"] == 5.0
  
imports: import typing
  from functools import partial
  import jax
  import jax.numpy
  import pytest
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling the TestGrad obj's underlying member method code with example usage: obj.test_grad() and returns: unknown