code: def test_grad(self):
      f = lambda x: x**2
      g = jax.grad(f)
      x = tf.constant(4.0)
      self.assertAllClose(g(x).numpy(), (2 * x).numpy())
  
file_path: sonnet/sonnet/src/functional/jax_test.py
imports: from absl.testing import parameterized
  from sonnet.src import test_utils
  from sonnet.src.functional import jax
  import tensorflow
label_desc: Use image node_id 1 for calling the GradTest obj's underlying member method code with example usage: obj.test_grad() without return types