code: def test_vmap_structured_in_axes(self):
      a, b, c, d = 2, 3, 4, 5
      k = 6  # batch size
      x = np.ones((k, a, b))  # batch axis in different locations
      y = np.ones((b, k, c))
      z = np.ones((c, d, k))
  
      def foo(tree_arg):
          x, (y, z) = tree_arg
          return tf_np.dot(x, tf_np.dot(y, z))
  
      tree = (x, (y, z))
      vfoo = extensions.vmap(foo, in_axes=((0, (1, 2)),))
      self.assertEqual(vfoo(tree).shape, (6, 2, 5))
  
      Point = collections.namedtuple("Point", ["x", "y"])
      tree = (x, Point(y, z))
      vfoo = extensions.vmap(foo, in_axes=((0, Point(1, 2)),))
      self.assertEqual(vfoo(tree).shape, (6, 2, 5))
  
      def foo2(tree_arg):
          x, dct = tree_arg
          y, z = dct["a"], dct["b"]
          return tf_np.dot(x, tf_np.dot(y, z))
  
      tree = (x, {"a": y, "b": z})
      vfoo = extensions.vmap(foo2, in_axes=((0, {"a": 1, "b": 2}),))
      self.assertEqual(vfoo(tree).shape, (6, 2, 5))
  
      tree = (x, collections.OrderedDict([("a", y), ("b", z)]))
      vfoo = extensions.vmap(
          foo2,
          in_axes=((0, collections.OrderedDict([("a", 1), ("b", 2)])),),
      )
      self.assertEqual(vfoo(tree).shape, (6, 2, 5))
  
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 7 for calling the VmapTest obj's underlying member method code with example usage: obj.test_vmap_structured_in_axes() and returns: tf_np, tf_np