code: def test_grad_with_type_predicate(self):
      m = nnx.Dict(
          a=nnx.Sequence([nnx.Param(10.0), nnx.BatchStat(20.0)]),
          b=nnx.Param(10.0),
          c=7,
          d=5.0,
      )
  
      @partial(nnx.grad, wrt=nnx.BatchStat)
      def f(m: nnx.Dict):
          # sum all params
          return m.a[0] + m.a[1] + m.b
  
      grads = f(m)
  
      assert isinstance(grads, nnx.State)
      assert grads["a"]["1"] == 1.0
      assert isinstance(grads.a.variables["1"], nnx.BatchStat)
      assert len(grads) == 1
  
      m.update(grads)
  
      assert m.a[0] == 10.0
      assert m.a[1] == 1.0
      assert m.b == 10.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 3 for calling the TestGrad obj's underlying member method code with example usage: obj.test_grad_with_type_predicate() and returns: unknown