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: Call the compiled function or module on multiple devices in parallel.
  Important: Make sure you call this function within the scope of VarCollection.replicate() statement.
docstr_tok: ["Call","the","compiled","function","or","module","on","multiple","devices","in","parallel",".","Important",":","Make","sure","you","call","this","function","within","the","scope","of","VarCollection.replicate","(",")","statement","."]
code: def __call__(self, *args):
      """Call the compiled function or module on multiple devices in parallel.
      Important: Make sure you call this function within the scope of VarCollection.replicate() statement.
      """
      unreplicated = [
          k
          for k, v in self.vc.items()
          if not isinstance(
              v.value,
              (
                  jax.Array,
                  jax.interpreters.partial_eval.JaxprTracer,
                  jax.interpreters.partial_eval.DynamicJaxprTracer,
              ),
          )
      ]
      assert not unreplicated, (
          f"Some variables were not replicated: {unreplicated}."
          "did you forget to call VarCollection.replicate on them?"
      )
  
      args = [
          x
          if i in self.static_argnums
          else jax.tree_map(self.device_reshape, [x])[0]
          for i, x in enumerate(args)
      ]
      output, changes = self._call(
          self.vc.tensors(),
          self.vc.subset(RandomState).tensors(),
          *args,
      )
      self.vc.assign(changes)
      return jax.tree_map(self.reduce, 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 3 for calling the Parallel obj's underlying member method code with example usage: obj.__call__() and returns: jax