put_along_axis

paddle. put_along_axis ( arr, indices, values, axis, reduce='assign', include_self=True, broadcast=True ) [source]

Put values into the destination array by given indices matrix along the designated axis.

Parameters
  • arr (Tensor) – The Destination Tensor. Supported data types are float32 and float64.

  • indices (Tensor) – Indices to put along each 1d slice of arr. This must match the dimension of arr, and need to broadcast against arr. Supported data type are int and int64.

  • axis (int) – The axis to put 1d slices along.

  • reduce (str, optional) – The reduce operation, default is ‘assign’, support ‘add’, ‘assign’, ‘mul’ and ‘multiply’.

  • include_self (bool, optional) – whether to reduce with the elements of arr. (Only support True now)

  • broadcast (bool, optional) – whether to broadcast indices.

Returns

Tensor, The indexed element, same dtype with arr

Examples

>>> import paddle

>>> x = paddle.to_tensor([[10, 30, 20], [60, 40, 50]])
>>> index = paddle.to_tensor([[0]])
>>> value = 99
>>> axis = 0
>>> result = paddle.put_along_axis(x, index, value, axis)
>>> print(result)
Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,
[[99, 99, 99],
 [60, 40, 50]])