import paddle
import torch
import numpy as np
# paddle
x_pd = paddle.randn([40, 50])
x_pd.stop_gradient = False
axis = 0
repeats_pd = paddle.randint(1, 50, [x_pd.shape[axis]])
y_pd = paddle.repeat_interleave(x_pd, repeats_pd, axis)
dy_pd = paddle.randn_like(y_pd)
dy_pd.stop_gradient = False
g_pd = paddle.grad(y_pd, x_pd, dy_pd, create_graph=True)[0]
ddx_pd = paddle.randn_like(x_pd)
gg_pd = paddle.grad(g_pd, dy_pd, ddx_pd)[0]
# torch
x_pt = torch.from_dlpack(x_pd.detach())
x_pt.requires_grad = True
repeats_pt = torch.from_dlpack(repeats_pd.detach())
y_pt = torch.repeat_interleave(x_pt, repeats_pt, axis)
np.testing.assert_allclose(y_pd.numpy(), y_pt.detach().cpu().numpy(), 1e-6, 1e-6)
dy_pt = torch.from_dlpack(dy_pd.detach())
dy_pt.requires_grad = True
g_pt = torch.autograd.grad(y_pt, x_pt, dy_pt, create_graph=True)[0]
np.testing.assert_allclose(g_pd.numpy(), g_pt.detach().cpu().numpy(), 1e-5, 1e-5)
ddx_pt = torch.from_dlpack(ddx_pd.detach())
gg_pt = torch.autograd.grad(g_pt, dy_pt, ddx_pt, create_graph=True)[0]
np.testing.assert_allclose(gg_pd.numpy(), gg_pt.detach().cpu().numpy(), 1e-5, 1e-5)