docstr: Vectorize module with a two arguments with mixed batch axis.
docstr_tok: ["Vectorize","module","with","a","two","arguments","with","mixed","batch","axis","."]
code: def test_vectorize_module_two_args_mixed_axis(self):
      """Vectorize module with a two arguments with mixed batch axis."""
      f1 = objax.nn.Linear(3, 4)
      f2 = objax.nn.Linear(5, 3)
      f = lambda x, y: jn.concatenate([f1(x), f2(y)], axis=1)
      fv = objax.Vectorize(
          lambda x, y: jn.concatenate([f1(x), f2(y)], axis=0),
          f1.vars("f1") + f2.vars("f2"),
          batch_axis=(0, 1),
      )
      x1 = objax.random.normal((96, 3))
      x2 = objax.random.normal((96, 5))
      y = f(x1, x2)
      yv = fv(x1, x2.T)
      self.assertTrue(jn.array_equal(y, yv))
  
file_path: objax/tests/vectorize.py
imports: import unittest
  import jax.numpy
  import numpy
  import objax
label_desc: Use image node_id 4 for calling the TestVectorize obj's underlying member method code with example usage: obj.test_vectorize_module_two_args_mixed_axis() without return types