code: def test_vmap_mismatched_axis_sizes_error_message_issue_705(self):
      # https://github.com/google/jax/issues/705
      with self.assertRaisesRegex(
          ValueError,
          "vmap must have at least one non-None value in in_axes",
      ):
          # If the output is mapped, there must be a non-None in_axes
          extensions.vmap(lambda x: x, in_axes=None)(
              tf_np.array([1.0, 2.0])
          )
  
      # Error is: TypeError: only integer scalar arrays can be converted to a
      # scalar index
      with self.assertRaisesRegex(
          ValueError,
          "vmap out_axes specification must be a tree prefix of the "
          "corresponding value.*",
      ):
          extensions.vmap(lambda x: x, in_axes=0, out_axes=(2, 3))(
              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 6 for calling the VmapTest obj's underlying member method code with example usage: obj.test_vmap_mismatched_axis_sizes_error_message_issue_705() without return types