mlx.core.slice_update

Contents

mlx.core.slice_update#

slice_update(a: array, update: array, start_indices: array, axes: Sequence[int], *, stream: None | Stream | Device = None) → array#

Update a sub-array of the input array.

Parameters:

  • a (array) – The input array to update

  • update (array) – The update array.

  • start_indices (array) – The index location to start the slice at.

  • axes (tuple(int)) – The axes corresponding to the indices in start_indices.

Returns:

The output array with the same shape and type as the input.

Return type:

array

Example

>>> a = mx.zeros((3, 3))
>>> mx.slice_update(a, mx.ones((1, 2)), start_indices=mx.array(1, 1), axes=(0, 1))
array([[0, 0, 0],
       [0, 1, 0],
       [0, 1, 0]], dtype=float32)
Contents