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: Vectorize constructor.
  
  Args:
      f: the function or the module to compile for vectorization.
      vc: the VarCollection of variables used by the function or module. This argument is required for functions.
      batch_axis: tuple of int or None for each of f's input arguments: the axis to use as batch during
          vectorization. Use None to automatically broadcast.
docstr_tok: ["Vectorize","constructor",".","Args",":","f",":","the","function","or","the","module","to","compile","for","vectorization",".","vc",":","the","VarCollection","of","variables","used","by","the","function","or","module",".","This","argument","is","required","for","functions",".","batch_axis",":","tuple","of","int","or","None","for","each","of","f","'s","input","arguments",":","the","axis","to","use","as","batch","during","vectorization",".","Use","None","to","automatically","broadcast","."]
code: def __init__(
      self,
      f: Union[Module, Callable],
      vc: Optional[VarCollection] = None,
      batch_axis: Tuple[Optional[int], ...] = (0,),
  ):
      """Vectorize constructor.
  
      Args:
          f: the function or the module to compile for vectorization.
          vc: the VarCollection of variables used by the function or module. This argument is required for functions.
          batch_axis: tuple of int or None for each of f's input arguments: the axis to use as batch during
              vectorization. Use None to automatically broadcast.
      """
      if not isinstance(f, Module):
          if vc is None:
              raise ValueError(
                  "You must supply the VarCollection used by the function f."
              )
          f = Function(f, vc)
  
      def vmap(
          tensor_list: List[JaxArray],
          random_list: List[JaxArray],
          *args,
      ):
          original_values = self.vc.tensors()
          try:
              self.vc.assign(tensor_list)
              self.vc.subset(RandomState).assign(random_list)
              return f(*args), self.vc.tensors()
          finally:
              self.vc.assign(original_values)
  
      fargs = positional_args_names(f)
      assert len(batch_axis) >= len(
          fargs
      ), f"The batched argument must be specified for all of {f} arguments {fargs}"
      self.batch_axis = batch_axis
      self.batch_axis_argnums = [
          (x, v) for x, v in enumerate(batch_axis) if v is not None
      ]
      assert (
          self.batch_axis_argnums
      ), f"No arguments to function {f} are vectorizable"
      self.vc = vc or f.vars()
      self._call = jax.vmap(vmap, (None, 0) + batch_axis)
      self.__wrapped__ = f
  
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 1 to create a new Vectorize object from inherited base classes: Module with example: obj = Vectorize(f, vc, batch_axis)