code: def test_negative_axes(self):
      x = np.arange(3 * 4 * 5).reshape(3, 4, 5)
      self.assertAllClose(
          extensions.vmap(tf_np.sum, in_axes=-3)(x),
          tf_np.sum(x, axis=(1, 2)),
      )
      self.assertAllClose(
          extensions.vmap(tf_np.sum, in_axes=-2)(x),
          tf_np.sum(x, axis=(0, 2)),
      )
      self.assertAllClose(
          extensions.vmap(tf_np.sum, in_axes=-1)(x),
          tf_np.sum(x, axis=(0, 1)),
      )
  
      identity = lambda y: y
      self.assertAllClose(
          x, extensions.vmap(identity, in_axes=0, out_axes=-3)(x)
      )
      self.assertAllClose(
          x.transpose(1, 0, 2),
          extensions.vmap(identity, in_axes=0, out_axes=-2)(x),
      )
      self.assertAllClose(
          x.transpose(1, 2, 0),
          extensions.vmap(identity, in_axes=0, out_axes=-1)(x),
      )
  
      self.assertAllClose(
          np.full((5,), 7),
          extensions.vmap(
              lambda *xs: xs, in_axes=(0, None), out_axes=(0, -1)
          )(np.arange(5), 7)[1],
      )
  
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 9 for calling the VmapTest obj's underlying member method code with example usage: obj.test_negative_axes() without return types