import paddle
x = paddle.randn([10, 20, 30, 40])
x.stop_gradient = False
axis = 2
indices = paddle.randint(0, x.shape[axis], [10, 20, 50, 40])
values = paddle.randn(indices.shape)
values.stop_gradient = False
out = paddle.put_along_axis(
    x,
    indices,
    values,
    axis,
    'add',
)
dout = paddle.randn_like(out)
dout.stop_gradient = False
dx, dv = paddle.grad(
    out,
    [x, values],
    dout,
    create_graph=True,
)
print(f"dx.shape = {dx.shape}, dv.shape = {dv.shape}")
ddx = paddle.randn_like(dx)
ddx.stop_gradient = False
ddv = paddle.randn_like(dv)
ddv.stop_gradient = False
ddout = paddle.grad(
    [dx, dv], # 此处可以仅保留 dx 或者 dv，以测试optional input分支
    dout,
    [ddx, ddv], # 此处可以仅保留 ddx 或者 ddv，以测试optional input分支
)[0]
print(f"ddout.shape = {ddout.shape}")
import torch
tx = torch.as_tensor(x.detach()).cuda()
tx.requires_grad = True
tindices = torch.as_tensor(indices.detach()).cuda()
tvalues = torch.as_tensor(values.detach()).cuda()
tvalues.requires_grad = True
tout = torch.scatter_add(
    tx,
    axis,
    tindices,
    tvalues,
)
tdout = torch.as_tensor(out.detach())
tdout.requires_grad = True
tdx, tdv = torch.autograd.grad(
    tout,
    [tx, tvalues],
    tdout,
    create_graph=True,
)
print(f"dx.shape = {dx.shape}, dv.shape = {dv.shape}")
tddx = torch.as_tensor(ddx.detach())
tddx.requires_grad = True
tddv = torch.as_tensor(ddv.detach())
tddv.requires_grad = True
tddout = torch.autograd.grad(
    [tdx, tdv], # 此处可以仅保留 tdx 或者 tdv，以测试optional input分支
    tdout,
    [tddx, tddv], # 此处可以仅保留 tddx 或者 tddv，以测试optional input分支
)[0]
print(f"ddout.shape = {ddout.shape}")
import numpy as np
np.testing.assert_allclose(ddout.numpy(), tddout.detach().cpu().numpy())