split_func_call#

ivy.split_func_call(func, inputs, mode, /, *, max_chunk_size=None, chunk_size=None, input_axes=0, output_axes=None, stop_gradients=False, device=None)[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 device the arrays are on.

Parameters:

func (Callable) – The function to be called.
inputs (Union[Array, NativeArray]) – A list of inputs to pass into the function.
mode (str) – The mode by which to unify the return values, must be one of [ concat | mean | sum ]
max_chunk_size (Optional[int], default: None) – The maximum size of each of the chunks to be fed into the function.
chunk_size (Optional[int], default: None) – The size of each of the chunks to be fed into the function. Specifying this arg overwrites the global split factor. Default is None.
input_axes (Union[int, Iterable[int]], default: 0) – The axes along which to split each of the inputs, before passing to the function. Default is 0.
output_axes (Optional[Union[int, Iterable[int]]], default: None) – The axes along which to concat each of the returned outputs. Default is same as fist input axis.
stop_gradients (bool, default: False) – Whether to stop the gradients for each computed return. Default is False.
device (Optional[Union[Device, NativeDevice]], default: None) – The device to set the split factor for. Sets the default device by default.

Return type:

Union[Array, NativeArray]

Returns:

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