code: def test_vmap_out_axes(self):
      f = extensions.vmap(lambda x: x, out_axes=0)
      inp = tf_np.arange(6).reshape([2, 3])
      self.assertAllClose(inp, f(inp))
      self.assertAllClose([inp, inp], f((inp, inp)))
  
      f = extensions.vmap(lambda x: x, out_axes=-1)
      self.assertAllClose(inp.T, f(inp))
  
      f = extensions.vmap(lambda x: x, out_axes=None)
      self.assertAllClose(inp[0], f(inp))
  
      f = extensions.vmap(
          lambda x: x, out_axes=([0], (-1, None), {"a": 1})
      )
      a, b, c = f(([inp], (inp, inp), {"a": inp}))
      self.assertAllClose([inp], a)
      self.assertAllClose((inp.T, inp[0]), b)
      self.assertAllClose(inp.T, c["a"])
  
file_path: trax/trax/tf_numpy/jax_tests/vmap_test.py
imports: import collections
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from trax.tf_numpy import extensions
  import trax.tf_numpy.numpy
  from tensorflow.python.ops.numpy_ops import np_math_ops
label_desc: Use image node_id 8 for calling the VmapTest obj's underlying member method code with example usage: obj.test_vmap_out_axes() without return types