[Prim&Comp] Support repeat_interleave double grad

[HydrogenSulfate]

PR Category

Operator Mechanism

PR Types

New features

Description

Pcard-75624

支持动态图和静态图下的repeat_interleave_with_tensor_index和repeat_interleave两个算子的二阶反向组合算子

相关PR：PaddlePaddle/PaddleScience#1160 @BeingGod

torch精度对比测试代码

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)

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.