[Prim] Support masked_select_double_grad in eager mode

[HydrogenSulfate]

PR Category

Operator Mechanism

PR Types

New features

Description

Pcard-75624

支持动态图下的 masked_select 二阶反向组合算子

torch精度对比测试代码

import paddle
import torch
import numpy as np


# paddle
x_pd = paddle.randn([1024, 10240])
x_pd.stop_gradient = False

mask_pd = paddle.randint(0, 2, x_pd.shape).astype("bool")

y_pd = paddle.masked_select(x_pd, mask_pd)

dy_pd = paddle.randn_like(y_pd)
dy_pd.stop_gradient = False

dx_pd, = paddle.grad(y_pd, x_pd, dy_pd, create_graph=True)
ddx_pd = paddle.randn_like(dx_pd)
ddx_pd.stop_gradient = False

ddy_pd, = paddle.grad(dx_pd, dy_pd, ddx_pd, create_graph=True)

# torch
x_pt = torch.from_dlpack(x_pd.detach()).requires_grad_(True)
mask_pt = torch.from_dlpack(mask_pd)
y_pt = torch.masked_select(x_pt, mask_pt)

dy_pt = torch.from_dlpack(dy_pd.detach()).requires_grad_(True)
dx_pt, = torch.autograd.grad(y_pt, x_pt, dy_pt, create_graph=True)

ddx_pt = torch.from_dlpack(ddx_pd.detach()).requires_grad_(True)
ddy_pt, = torch.autograd.grad(dx_pt, dy_pt, ddx_pt, create_graph=True)

np.testing.assert_allclose(y_pd.numpy(), y_pt.detach().cpu().numpy())
np.testing.assert_allclose(dx_pd.numpy(), dx_pt.detach().cpu().numpy())
np.testing.assert_allclose(ddy_pd.numpy(), ddy_pt.detach().cpu().numpy())