code: def test_vmap_in_axes_list(self):
      # https://github.com/google/jax/issues/2367
      dictionary = {"a": 5.0, "b": tf_np.ones(2)}
      x = tf_np.zeros(3)
      y = tf_np.arange(3.0)
  
      def f(dct, x, y):
          return dct["a"] + dct["b"] + x + y
  
      out1 = extensions.vmap(f, (None, 0, 0))(dictionary, x, y)
      out2 = extensions.vmap(f, [None, 0, 0])(dictionary, x, y)
      self.assertAllClose(out1, out2)
  
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 1 for calling the VmapTest obj's underlying member method code with example usage: obj.test_vmap_in_axes_list() and returns: unknown