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: Utility to reshape an input array in order to broadcast to multiple devices.
docstr_tok: ["Utility","to","reshape","an","input","array","in","order","to","broadcast","to","multiple","devices","."]
code: def device_reshape(self, x: JaxArray) -> JaxArray:
      """Utility to reshape an input array in order to broadcast to multiple devices."""
      assert hasattr(x, "ndim"), (
          f"Expected JaxArray, got {type(x)}. If you are trying to pass a scalar to "
          f"parallel, first convert it to a JaxArray, for example np.float(0.5)"
      )
      if x.ndim == 0:
          return np.broadcast_to(x, [self.ndevices])
      assert x.shape[0] % self.ndevices == 0, (
          f"Must be able to equally divide batch {x.shape} among "
          f"{self.ndevices} devices, but does not go equally."
      )
      return x.reshape(
          (self.ndevices, x.shape[0] // self.ndevices) + x.shape[1:]
      )
  
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 Parallel obj's underlying member method code with example usage: obj.device_reshape(x) and returns: x, np