Multi node

Collection of multi-node Ivy functions.

class ivy.ClonedItem(data: Dict[ivy.Node, Any], axis=0)[source]

Bases: ivy.core.multi_node.MultiNodeItem

class ivy.ClonedIter(data: Iterable, node_strs)[source]

Bases: ivy.core.multi_node.MultiNodeIter

class ivy.ClonedNest(data: Iterable, node_strs, max_depth=1)[source]

Bases: ivy.core.multi_node.MultiNodeNest

class ivy.MultiNode(data: Iterable, axis=0)[source]

Bases: object

__init__(data: Iterable, axis=0)[source]

Initialize self. See help(type(self)) for accurate signature.

class ivy.MultiNodeItem(data: Dict[ivy.Node, Any], axis=0)[source]

Bases: ivy.core.multi_node.MultiNode

__init__(data: Dict[ivy.Node, Any], axis=0)[source]

Initialize self. See help(type(self)) for accurate signature.

items()[source]
keys()[source]
property shape
values()[source]
class ivy.MultiNodeIter(data: Iterable, node_strs)[source]

Bases: ivy.core.multi_node.MultiNode

__init__(data: Iterable, node_strs)[source]

Initialize self. See help(type(self)) for accurate signature.

at_node(node_str)[source]
at_nodes()[source]
class ivy.MultiNodeNest(data: Iterable, node_strs, max_depth=1)[source]

Bases: ivy.core.multi_node.MultiNodeIter

__init__(data: Iterable, node_strs, max_depth=1)[source]

Initialize self. See help(type(self)) for accurate signature.

at_node(node_str)[source]
class ivy.NodeDistItem(data: Dict[ivy.Node, Any], axis=0)[source]

Bases: ivy.core.multi_node.MultiNodeItem

class ivy.NodeDistIter(data: Iterable, node_strs)[source]

Bases: ivy.core.multi_node.MultiNodeIter

class ivy.NodeDistNest(data: Iterable, node_strs, max_depth=1)[source]

Bases: ivy.core.multi_node.MultiNodeNest

class ivy.NodeManager(node_mapper=None, node_strs: Optional[Union[Iterable[str], Dict[str, int]]] = None, da_dim_size=None, safety_factor=1.1, min_node_dim_size=0, max_node_dim_step_ratio=0.1, min_unit_node_tune_steps=10, min_sf_tune_steps=10, starting_split_factor=0.0, max_split_factor_step_size=0.05, tune_node_alloc=True, tune_node_splits=True)[source]

Bases: object

__init__(node_mapper=None, node_strs: Optional[Union[Iterable[str], Dict[str, int]]] = None, da_dim_size=None, safety_factor=1.1, min_node_dim_size=0, max_node_dim_step_ratio=0.1, min_unit_node_tune_steps=10, min_sf_tune_steps=10, starting_split_factor=0.0, max_split_factor_step_size=0.05, tune_node_alloc=True, tune_node_splits=True)[source]

Create node manager, which unlike the node mapper, handles all argument cloning and distributing internally. The node manager only receivess a specification regarding the ratio of the batch each node should consume.

Parameters
  • node_mapper (NodeMapper) – The pre-built node mapper used by the manager internally.

  • node_strs (sequence of strs or dict of split sizes) – The nodes to distribute and clone the arguments across.

  • da_dim_size (int) – The size of the dimension along which the node allocation splitting is performed.

  • safety_factor (float, optional) – The factor by which to be safe in the avoidance of OOM GPU errors. Default is 1.1.

  • min_node_dim_size (int, optional) – The minimum dimension size to pass to a node. Default is 0.

  • max_node_dim_step_ratio (int, optional) – The maximum step ratio for changing the dimension for a node. Default is 0.1.

  • min_unit_node_tune_steps (int, optional) – The minimum number of tune steps to make when optimizing with unit step size. Default is 10.

  • min_sf_tune_steps (int, optional) – Minimum number of split factor tune steps. Default is 10.

  • starting_split_factor (float, optional) – The initial node-specific split factor. Default is 0.

  • max_split_factor_step_size (float, optional) – The maximum step size for changing the split factor for a node. Default is 0.05.

  • tune_node_alloc (bool, optional) – Whether to tune the node split sizes internally based on node utilization tracking, and use the provided values for initialization. Default is True.

  • tune_node_splits (bool, optional) – Whether to tune the per-node split sizes internally. Default is True.

da_tune_step()[source]
property dim_size
map(cloned=None, to_clone=None, distributed=None, to_distribute=None)[source]

Map the function fn to each of the MultiNode args and kwargs, running each function in parallel with CUDA-safe multiprocessing.

Parameters
  • cloned (dict of any, optional) – The MutliNode keyword arguments which are already cloned. Default is None.

  • to_clone (dict of any, optional) – The MutliNode keyword arguments to clone and map to the function. Default is None.

  • distributed (dict of any, optional) – The MutliNode keyword arguments which already distributed. Default is None.

  • to_distribute (dict of any, optional) – The MutliNode keyword arguments to distribute and map to the function. Default is None.

Returns

The results of the function, returned as a MultiNode instance.

ns_tune_step()[source]
repeated_config_check()[source]
property tune_step
class ivy.NodeMapper(fn, ret_fn, queue_class, worker_class, node_strs, timeout=None, constant=None, unique=None)[source]

Bases: abc.ABC

__init__(fn, ret_fn, queue_class, worker_class, node_strs, timeout=None, constant=None, unique=None)[source]

Node Mapper base class.

Parameters
  • fn (callable) – The function which the node mapper parallelises across nodes.

  • ret_fn (callable) – The function which receives the ivy.MultiNodeIter as input, and produces a single node output.

  • queue_class (class) – The class to use for creating queues.

  • worker_class (class) – The class to use for creating parallel workers.

  • node_strs (sequence of str) – A list of nodes on which to parallelise the function.

  • timeout (float, optional) – The timeout for getting items from the queues. Default is global.

  • constant (dict of any, optional) – A dict of keyword arguments which are the same for each process. Default is None.

  • unique (dict of iterables of any, optional) – A dict of keyword argument sequences which are unique for each process. Default is None.

map(used_node_strs=None, split_factors=None, **kwargs)[source]

Map the function fn to each of the MultiNode args and kwargs, running each function in parallel with CUDA-safe multiprocessing.

Parameters
  • used_node_strs (sequence of str, optional) – The nodes used in the current mapping pass. Default is all node_strs.

  • split_factors (dict of floats, optional) – The updated split factors 0 < sf < 1 for each node. Default is None.

  • kwargs (dict of any) – The MultiNode keyword arguments to map the function to.

Returns

The results of the function, returned as a MultiNode instance.

class ivy.NodeMapperMultiProc(fn, ret_fn, node_strs, timeout=None, constant=None, unique=None)[source]

Bases: ivy.core.multi_node.NodeMapper

__init__(fn, ret_fn, node_strs, timeout=None, constant=None, unique=None)[source]

Node Mapper base class.

Parameters
  • fn (callable) – The function which the node mapper parallelises across nodes.

  • ret_fn (callable) – The function which receives the ivy.MultiNodeIter as input, and produces a single node output.

  • queue_class (class) – The class to use for creating queues.

  • worker_class (class) – The class to use for creating parallel workers.

  • node_strs (sequence of str) – A list of nodes on which to parallelise the function.

  • timeout (float, optional) – The timeout for getting items from the queues. Default is global.

  • constant (dict of any, optional) – A dict of keyword arguments which are the same for each process. Default is None.

  • unique (dict of iterables of any, optional) – A dict of keyword argument sequences which are unique for each process. Default is None.

class ivy.Profiler(save_dir)[source]

Bases: abc.ABC

__init__(save_dir)[source]

Initialize self. See help(type(self)) for accurate signature.

abstract start()[source]
abstract stop()[source]
ivy.clone(x, node_strs)[source]

Clone the input item to each of the specified nodes, returning a list of cloned items, each on a different node.

Parameters
  • x (array or container) – The input array or container to clone to each node.

  • node_strs (sequence of strs) – The nodes to clone the input to.

Returns

array or container distributed across the target nodes

ivy.clone_array(x, node_strs)[source]

Clone an array across the specified nodes, returning a list of cloned arrays, each on a different node.

Parameters
  • x (array) – The array to clone across nodes.

  • node_strs (sequence of strs) – The nodes to clone the array to.

Returns

array cloned to each of the target nodes

ivy.clone_iter(xs, node_strs)[source]

Clone elements of the iterbale xs to each of the specified nodes.

Parameters
  • xs (iterable of any) – The iterable of items to clone.

  • node_strs (sequence of strs) – The nodes to clone each of the iterable elements to.

Returns

iterable with each element cloned to each of the target nodes

ivy.clone_nest(args, kwargs, node_strs, max_depth=1)[source]

Clone the input arguments across the specified nodes.

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

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

  • node_strs (sequence of strs) – The nodes to clone the arguments to.

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

Returns

arguments cloned to each of the target nodes

ivy.default_node()[source]

Return the default node.

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

Distribute the nested input arguments across the specified nodes.

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

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

  • node_strs (sequence of strs or dict of split sizes) – The nodes 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 nodes

ivy.gpu_is_available(f: Optional[ivy.Framework] = None) → bool[source]

Determine whether a GPU is available to use, with the backend framework.

Parameters

f (ml_framework, optional) – Machine learning framework. Inferred from inputs if None.

Returns

Boolean, as to whether a gpu is available.

ivy.node(x: Union[ivy.Array, ivy.NativeArray], f: Optional[ivy.Framework] = None) → ivy.Node[source]

Get the native node handle for input array x.

Parameters
  • x (array) – Tensor for which to get the node handle.

  • f (ml_framework, optional) – Machine learning framework. Inferred from inputs if None.

Returns

node handle for the array, in native framework format.

ivy.node_dist(x, node_strs: Union[Iterable[str], Dict[str, int]], axis=0)[source]

Distribute the input item across the specified nodes, returning a list of sub-items, each on a different node.

Parameters
  • x (array or container) – The input array or container to distribute across nodes.

  • node_strs (sequence of strs or dict of split sizes) – The nodes to distribute the input across.

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

Returns

array or container distributed across the target nodes

ivy.node_dist_array(x, node_strs: Union[Iterable[str], Dict[str, int]], axis=0)[source]

Distribute an array across the specified nodes, returning a list of sub-arrays, each on a different node.

Parameters
  • x (array) – The array to distribute across nodes.

  • node_strs (sequence of strs or dict of split sizes) – The nodes to distribute the array across.

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

Returns

array distributed across the target nodes

ivy.node_dist_iter(xs, node_strs: Union[Iterable[str], Dict[str, int]], axis=0)[source]

Distribute elements of the iterbale xs across the specified nodes.

Parameters
  • xs (iterable of any) – The iterable of items to distribute.

  • node_strs (sequence of strs or dict of split sizes) – The nodes to distribute the iterable elements across.

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

Returns

iterable with each element distributed to the target nodes

ivy.node_str(x: Union[ivy.Array, ivy.NativeArray], f: Optional[ivy.Framework] = None) → str[source]

Get the node string for input array x.

Parameters
  • x (array) – Tensor for which to get the node string.

  • f (ml_framework, optional) – Machine learning framework. Inferred from inputs if None.

Returns

node string for the array, e.g. ‘cuda:0’, ‘cuda:1’, ‘cpu’ etc..

ivy.node_to_str(node_in: ivy.Node, f: Optional[ivy.Framework] = None) → str[source]

Convert native data type to string representation.

Parameters
  • node_in (node handle) – The node handle to convert to string.

  • f (ml_framework, optional) – Machine learning framework. Inferred from inputs if None.

Returns

node string e.g. ‘cuda:0’.

ivy.node_util(node_str: str) → float[source]

Get the current utilization (%) for a given node.

Parameters

node_str (str) – The node string of the node to query utilization for.

Returns

The node utilization (%)

ivy.num_gpus(f: Optional[ivy.Framework] = None) → int[source]

Determine the number of available GPUs, with the backend framework.

Parameters

f (ml_framework, optional) – Machine learning framework. Inferred from inputs if None.

Returns

Number of available GPUs.

ivy.percent_used_mem_on_node(node_str: str, process_specific=False) → float[source]

Get the percentage used memory for a given node string. In case of CPU, the used RAM is returned.

Parameters
  • node_str (str) – The node string to conver to native node handle.

  • process_specific (bool, optional) – Whether the check the memory used by this python process alone. Default is False.

Returns

The percentage used memory on the node.

ivy.set_default_node(node)[source]
ivy.set_split_factor(factor, node_str=None)[source]

Set the global split factor for a given node, which can be used to scale batch splitting chunk sizes for the node across the codebase.

Parameters
  • factor (float) – The factor to set the node-specific split factor to.

  • node_str (str, optional) – The node to set the split factor for. Sets the default node by default.

ivy.split_factor(node_str=None)[source]

Get the global split factor for a given node, which can be used to scale batch splitting chunk sizes for the node across the codebase. Default global value for each node is 1.

Parameters

node_str (str, optional) – The node to query the split factor for. Sets the default node by default.

Returns

The split factor for the specified node.

ivy.split_func_call(func: Callable, inputs: Iterable[Union[ivy.Array, ivy.NativeArray, ivy.Container]], mode: str, max_chunk_size: Optional[int] = None, chunk_size: Optional[int] = None, input_axes: Union[int, Iterable[int]] = 0, output_axes: Optional[Union[int, Iterable[int]]] = None) → Iterable[Union[ivy.Array, ivy.NativeArray, ivy.Container]][source]

Call a function by splitting its inputs along a given axis, and calling the function in chunks, rather than feeding the entire input array at once. This can be useful to reduce memory usage of the node the arrays are on. :param func: The function to be called. :type func: callable :param inputs: A list of inputs to pass into the function. :type inputs: sequence of arrays :param mode: The mode by which to unify the return values, must be one of [ concat | mean | sum ] :type mode: str :param max_chunk_size: The maximum size of each of the chunks to be fed into the function. :type max_chunk_size: int :param chunk_size: The size of each of the chunks to be fed into the function. Specifying this arg overwrites the

global split factor. Default is None.

Parameters
  • input_axes (int or sequence of ints, optional) – The axes along which to split each of the inputs, before passing to the function. Default is 0.

  • output_axes (int or sequence of ints, optional) – The axes along which to concat each of the returned outputs. Default is same as fist input axis.

Returns

The return from the function, following input splitting and re-concattenation.

ivy.str_to_node(node_str: str, f: Optional[ivy.Framework] = None) → ivy.Node[source]

Convert node string representation to native node type.

Parameters
  • node_str (str) – The node string to conver to native node handle.

  • f (ml_framework, optional) – Machine learning framework. Inferred from inputs if None.

Returns

Native node handle.

ivy.to_node(x: Union[ivy.Array, ivy.NativeArray], node_str: Optional[str] = None, f: Optional[ivy.Framework] = None) → Union[ivy.Array, ivy.NativeArray][source]

Move the input array x to the desired node, specified by node string.

Parameters
  • x (array) – Array to move onto the node.

  • node_str (str, optional) – node to move the array to ‘cuda:0’, ‘cuda:1’, ‘cpu’ etc. Keep same node if None.

  • f (ml_framework, optional) – Machine learning framework. Inferred from inputs if None.

Returns

The array x, but now placed on the target node.

ivy.total_mem_on_node(node_str: str) → float[source]

Get the total amount of memory (in GB) for a given node string. In case of CPU, the total RAM is returned.

Parameters

node_str (str) – The node string to conver to native node handle.

Returns

The total memory on the node in GB.

ivy.tpu_is_available(f: Optional[ivy.Framework] = None) → bool[source]

Determine whether a TPU is available to use, with the backend framework.

Parameters

f (ml_framework, optional) – Machine learning framework. Inferred from inputs if None.

Returns

Boolean, as to whether a tpu is available.

ivy.unify(xs, node_str, mode, axis=0)[source]

Unify a list of sub-arrays, on arbitrary nodes, to a single concattenated array on the specified node.

Parameters
  • xs (sequence of arrays) – The list of sub-arrays to unify onto the specified node.

  • node_str (str) – The node to unify the sub-arrays 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 array, if concat mode is set. Default is 0.

Returns

array unified to the target node

ivy.unify_array(xs, node_str, mode, axis=0)[source]

Unify a list of sub-arrays, on arbitrary nodes, to a single array on the specified node.

Parameters
  • xs (sequence of arrays) – The list of arrays to unify onto the specified node.

  • node_str (str) – The node to unify the arrays 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 array, if concat mode is set. Default is 0.

Returns

array unified to the target node

ivy.unify_iter(xs, node_str, mode, axis=0, transpose=False)[source]

Unify elements of the iterbale xs to a single target node.

Parameters
  • xs (iterable of any) – The iterable of items to unify.

  • node_str (str) – The node to unify the elements of the iterable 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.

  • transpose (bool, optional) – Whether to transpose the first and second dimensions of the iterator. Default is False.

Returns

iterable with each element unified to a single target nodes

ivy.unify_nest(args: Type[ivy.core.multi_node.MultiNode], kwargs: Type[ivy.core.multi_node.MultiNode], node_str, mode, axis=0, max_depth=1)[source]

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

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

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

  • node_str (str) – The node 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 node

ivy.used_mem_on_node(node_str: str, process_specific=False) → float[source]

Get the used memory (in GB) for a given node string. In case of CPU, the used RAM is returned.

Parameters
  • node_str (str) – The node string to conver to native node handle.

  • process_specific (bool, optional) – Whether the check the memory used by this python process alone. Default is False.

Returns

The used memory on the node in GB.