dev_unify_nest

ivy.dev_unify_nest(args: Type[ivy.core.device.MultiDev], kwargs: Type[ivy.core.device.MultiDev], dev_str, mode, axis=0, max_depth=1)[source]

Unify the input nested arguments, which consist of sub-arrays spread across arbitrary devices, to unified arrays on the single target device.

Parameters
  • args (MultiDev) – The nested positional arguments to unify.

  • kwargs (MultiDev) – The nested keyword arguments to unify.

  • dev_str (str) – The device to unify the nested arguments to.

  • mode (str) – The mode by which to unify, must be one of [ concat | mean | sum ]

  • axis (int, optional) – The axis along which to concattenate the sub-arrays. 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 unified to the target device


Supported Frameworks:

empty jax_logo empty tf_logo empty pytorch_logo empty mxnet_logo empty numpy_logo empty