code: def test_vmap_in_axes_leaf_types(self):
      with self.assertRaisesRegex(
          TypeError, r"vmap in_axes must be an int, None, or .*"
      ):
          extensions.vmap(
              lambda x: x, in_axes=(tf_np.array([1.0, 2.0]),)
          )(tf_np.array([1.0, 2.0]))
  
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 3 for calling the VmapTest obj's underlying member method code with example usage: obj.test_vmap_in_axes_leaf_types() without return types