code: def testGradNonScalarOutput(self):
      def f(a):
          return a
  
      g = extensions.grad(f)
      with self.assertRaisesWithPredicateMatch(
          ValueError, r"result .* must be a scalar"
      ):
          g(tf_np.asarray([1.0, 2.0]))
  
      @extensions.jit
      def g_jitted(a):
          return extensions.grad(f)(a)
  
      g_jitted(tf_np.asarray(1.0))
      with self.assertRaisesWithPredicateMatch(
          ValueError, r"result .* must be a scalar"
      ):
          g_jitted(tf_np.asarray([1.0, 2.0]))
  
imports: import functools
  import itertools
  from absl import flags
  from absl.testing import parameterized
  from jax import lax
  import numpy
  import tensorflow.compat.v2
  from trax.tf_numpy import extensions
  import trax.tf_numpy.numpy
label_desc: Use image node_id 8 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testGradNonScalarOutput() and returns: a, extensions