class_docstr: Parallel module takes a function or a module and compiles it for running on multiple devices in parallel.
class_docstr_tok: ["Parallel","module","takes","a","function","or","a","module","and","compiles","it","for","running","on","multiple","devices","in","parallel","."]
docstr: Parallel constructor.
  
  Args:
      f: the function or the module to compile for parallelism.
      vc: the VarCollection of variables used by the function or module. This argument is required for functions.
      reduce: the function used reduce the outputs from many devices to a single device value.
      axis_name: what name to give to the device dimension, used in conjunction with objax.functional.parallel.
      static_argnums: tuple of indexes of f's input arguments to treat as static (constants)).
          A new graph is compiled for each different combination of values for such inputs.
docstr_tok: ["Parallel","constructor",".","Args",":","f",":","the","function","or","the","module","to","compile","for","parallelism",".","vc",":","the","VarCollection","of","variables","used","by","the","function","or","module",".","This","argument","is","required","for","functions",".","reduce",":","the","function","used","reduce","the","outputs","from","many","devices","to","a","single","device","value",".","axis_name",":","what","name","to","give","to","the","device","dimension",",","used","in","conjunction","with","objax.functional.parallel",".","static_argnums",":","tuple","of","indexes","of","f","'s","input","arguments","to","treat","as","static","(","constants",")",")",".","A","new","graph","is","compiled","for","each","different","combination","of","values","for","such","inputs","."]
code: def __init__(
      self,
      f: Union[Module, Callable],
      vc: Optional[VarCollection] = None,
      reduce: Callable[[JaxArray], JaxArray] = jn.concatenate,
      axis_name: str = "device",
      static_argnums: Optional[Tuple[int, ...]] = None,
  ):
      """Parallel constructor.
  
      Args:
          f: the function or the module to compile for parallelism.
          vc: the VarCollection of variables used by the function or module. This argument is required for functions.
          reduce: the function used reduce the outputs from many devices to a single device value.
          axis_name: what name to give to the device dimension, used in conjunction with objax.functional.parallel.
          static_argnums: tuple of indexes of f's input arguments to treat as static (constants)).
              A new graph is compiled for each different combination of values for such inputs.
      """
      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 pmap(
          tensor_list: List[jax.Array],
          random_list: List[jax.Array],
          *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)
  
      static_argnums = sorted(static_argnums or ())
      self.axis_name = axis_name
      self.ndevices = jax.local_device_count()
      self.reduce = reduce
      self.static_argnums = frozenset(static_argnums)
      self.vc = vc or f.vars()
      self._call = jax.pmap(
          pmap,
          axis_name=axis_name,
          static_broadcasted_argnums=[x + 2 for x in static_argnums],
      )
      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 Parallel object from inherited base classes: Module with example: obj = Parallel(f, vc, reduce, axis_name, static_argnums)