code: def test_argnums(self):
      f = lambda x, y: (x**2 + y**2)
      g = jax.grad(f, argnums=(0, 1))
      x = tf.constant(4.0)
      y = tf.constant(5.0)
      gx, gy = g(x, y)
      self.assertAllClose(gx.numpy(), (2 * x).numpy())
      self.assertAllClose(gy.numpy(), (2 * y).numpy(), rtol=1e-3)
  
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 2 for calling the GradTest obj's underlying member method code with example usage: obj.test_argnums() without return types