import paddle
def scatter_reduce(input: paddle.Tensor, 
                   axis: int, 
                   index: paddle.Tensor, 
                   src: paddle.Tensor, 
                   reduce: str) -> paddle.Tensor:
    # reduce "sum", "prod", "mean", "amax", "amin"
    if reduce == "sum":
        input.put_along_axis_(indices=index, values=src, axis=axis, reduce="add")
    elif reduce == "mean":
        input.put_along_axis_(indices=index, values=src, axis=axis, reduce="add")
        dst_div = (
            paddle.ones_like(input)
            .put_along_axis(
                indices=index, values=paddle.to_tensor(1.0, dtype=input.dtype), axis=axis, reduce="add"
            )
        )
        input = input / dst_div
    elif reduce == "prod":
        input = input.put_along_axis(indices=index, values=src, axis=axis, reduce="mul")
    else:
        raise NotImplementedError("only support mode in ['sum', 'prod', 'mean']!")
    return input
paddle.set_device('cpu')
src = paddle.to_tensor([1., 2., 3., 4., 5., 6.])
index = paddle.to_tensor([0, 1, 0, 1, 2, 1])
input = paddle.to_tensor([1., 2., 3., 4.])
out = scatter_reduce(input, 0, index, src, reduce="prod")
print(out)
# Tensor(shape=[4], dtype=float32, place=Place(cpu), stop_gradient=True,
#        [3. , 96., 15., 4. ])

paddle.set_device('gpu')
src = paddle.to_tensor([1., 2., 3., 4., 5., 6.])
index = paddle.to_tensor([0, 1, 0, 1, 2, 1])
input = paddle.to_tensor([1., 2., 3., 4.])
out = scatter_reduce(input, 0, index, src, reduce="prod")
print(out)
# Tensor(shape=[4], dtype=float32, place=Place(gpu:0), stop_gradient=True,
#        [1. , 4. , 15., 4. ])