code: def test_vmap_in_axes_tree_prefix_error(self):
      # https://github.com/google/jax/issues/795
      self.assertRaisesRegex(
          ValueError,
          "vmap in_axes specification must be a tree prefix of the corresponding "
          r"value, got specification \(0, 0\) for value tree ",
          lambda: extensions.vmap(lambda x: x, in_axes=(0, 0))(
              tf_np.ones(3)
          ),
      )
  
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 2 for calling the VmapTest obj's underlying member method code with example usage: obj.test_vmap_in_axes_tree_prefix_error() without return types