docstr: Test for custom_grad.
docstr_tok: ["Test","for","custom_grad","."]
code: def testCustomGrad(self):
      """Test for custom_grad."""
      x_shape = (tf.TensorShape([10]), tf.TensorShape([1, 10]))
      y_shape = tf.TensorShape([])
      dtype = np.float32
      scale1 = 5.0
      scale2 = 6.0
  
      def fwd(a, b):
          return tf_np.sum(tf_np.sqrt(tf_np.exp(a)) + b)
  
      @extensions.custom_grad
      def f(a, b):
          y = fwd(a, b)
  
          def vjp(dy):
              return dy * scale1 * a, dy * scale2 * b
  
          return y, vjp
  
      rng = tf.random.Generator.from_seed(1234)
      x, dy = tf.nest.map_structure(
          lambda shape: uniform(rng, shape, dtype), [x_shape, y_shape]
      )
      expected_y = fwd(*x)
      expected_dx = (dy * scale1 * x[0], dy * scale2 * x[1])
      y, vjp = extensions.vjp(f, *x)
      dx = vjp(dy)
      self.assertAllClose(expected_y, y)
      self.assertAllClose(expected_dx, dx)
  
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 3 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testCustomGrad() and returns: tf_np, y,  vjp, unknown,  unknown