dev_dist_nest

ivy.dev_dist_nest(args, kwargs, dev_strs: Union[Iterable[str], Dict[str, int]], axis=0, max_depth=1)[source]

Distribute the nested input arguments across the specified devices.

Parameters
  • args (list of any) – The positional nested arguments to distribute.

  • kwargs (dict of any) – The keyword nested arguments to distribute.

  • dev_strs (sequence of strs or dict of split sizes) – The devices to distribute the nested arguments across.

  • axis (int, optional) – The axis along which to split the arrays in the arguments. Default is 0.

  • max_depth (int, optional) – The maximum nested depth to reach. Default is 1. Increase this if the nest is deeper.

Returns

nested arguments distributed to the target devices


Supported Frameworks:

empty jax_logo empty tf_logo empty pytorch_logo empty mxnet_logo empty numpy_logo empty