code: def testVmap(self):
      fn1 = extensions.vmap(lambda z: z * z)
  
      x = tf_np.arange(10)
      self.assertAllClose(x * x, fn1(x))
  
      y = tf.range(10)
      np_y = tf_np.asarray(y)
      output = fn1(y)
      self.assertIsInstance(output, tf_np.ndarray)
      self.assertAllClose(np_y * np_y, output)
  
      fn2 = extensions.vmap(lambda x, y: x + y)
      x = tf_np.random.randn(10, 3)
      y = tf_np.random.randn(10, 2, 3)
      self.assertAllClose(tf_np.expand_dims(x, 1) + y, fn2(x, y))
  
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 43 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testVmap() without return types