mlx.core.put_along_axis

Contents

mlx.core.put_along_axis#

put_along_axis(a: array, /, indices: array, values: array, axis: int | None = None, *, stream: None | Stream | Device = None) array#

Put values along an axis at the specified indices.

Parameters:
  • a (array) – Destination array.

  • indices (array) – Indices array. These should be broadcastable with the input array excluding the axis dimension.

  • values (array) – Values array. These should be broadcastable with the indices.

  • axis (int or None) – Axis in the destination to put the values to. If axis == None the destination is flattened prior to the put operation.

Returns:

The output array.

Return type:

array