code: def test_has_aux(self):
      f = lambda x: (x**2, "aux")
      g = jax.grad(f, has_aux=True)
      x = tf.constant(2.0)
      gx, aux = g(x)
      self.assertAllClose(gx.numpy(), (2 * x).numpy())
      self.assertEqual(aux, "aux")
  
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 3 for calling the GradTest obj's underlying member method code with example usage: obj.test_has_aux() without return types