index_copy

→Tensor

index_copy(input: Tensor, dim: int, index: Tensor, source: Tensor)

source edit

Return a copy of input with slices at index replaced by source.

Single set-scatter: the 1-D index is broadcast along dim to source's shape and applied through the engine scatter_set primitive — one MPSGraph scatterAlongAxis (Set mode) op when compiled, or mlx.put_along_axis eager on GPU. Differentiable through both input and source.

Parameters

inputTensor

Destination tensor; not mutated (a fresh copy is returned).

dimint

Axis along which slices are addressed.

indexTensor

1-D int32 / int64 tensor of positions along dim. Length must equal source.shape[dim].

sourceTensor

Replacement slices. All non-dim dimensions must match input; source.shape[dim] must equal index.shape[0].

Returns

Tensor

Same shape and dtype as input; values at the indexed positions are taken from source, others from input.