code: def test_vmap_unbatched_object_passthrough_issue_183(self):
      # https://github.com/google/jax/issues/183
      fun = lambda f, x: f(x)
      vfun = extensions.vmap(fun, (None, 0))
      ans = vfun(lambda x: x + 1, tf_np.arange(3))
      self.assertAllClose(ans, np.arange(1, 4))
  
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 5 for calling the VmapTest obj's underlying member method code with example usage: obj.test_vmap_unbatched_object_passthrough_issue_183() without return types