class_docstr: Vectorize module takes a function or a module and compiles it for running in parallel on a single device.
class_docstr_tok: ["Vectorize","module","takes","a","function","or","a","module","and","compiles","it","for","running","in","parallel","on","a","single","device","."]
docstr: Call the vectorized version of the function or module.
docstr_tok: ["Call","the","vectorized","version","of","the","function","or","module","."]
code: def __call__(self, *args):
      """Call the vectorized version of the function or module."""
      assert len(args) == len(self.batch_axis), (
          f"Number of arguments passed {len(args)} must match "
          f"batched {len(self.batch_axis)}"
      )
      nsplits = args[self.batch_axis_argnums[0][0]].shape[
          self.batch_axis_argnums[0][1]
      ]
      output, changes = self._call(
          self.vc.tensors(),
          [v.split(nsplits) for v in self.vc.subset(RandomState)],
          *args,
      )
      for v, u in zip(self.vc, changes):
          v.reduce(u)
      return output
  
imports: from collections import namedtuple
  from typing import Optional, List, Union, Callable, Tuple
  import jax
  import jax.numpy
  import numpy
  from objax.typing import JaxArray
  from objax.util import class_name, override_args_kwargs, positional_args_names, repr_function
  from objax.variable import BaseVar, RandomState, VarCollection
label_desc: Use image node_id 2 for calling the Vectorize obj's underlying member method code with example usage: obj.__call__() and returns: output