[Prim] Support masked_select_grad for dynamic shape

[HydrogenSulfate]

PR Category

Operator Mechanism

PR Types

New features

Description

Pcard-75624

为masked_select_grad静态图组合算子支持动态shape

与torch对比精度验证代码

import paddle from paddle.static import InputSpec import torch import numpy as np paddle.framework.core._set_prim_all_enabled(True) for i in range(50): # paddle x_pd = paddle.randn([1024, 13, 14]) x_pd.stop_gradient = False mask_pd = paddle.randint(low=0, high=2, shape=[7, 1, 13, 1]).astype("bool") with paddle.no_grad(): t = paddle.masked_select(x_pd, mask_pd) dy_pd = paddle.randn(t.shape) dy_pd.stop_gradient = False ddx_pd = paddle.randn(x_pd.shape) ddx_pd.stop_gradient = False def g(x_pd_, mask_pd_, dy_pd, ddx_pd): y_pd = paddle.masked_select(x_pd_, mask_pd_) dx_pd, = paddle.grad(y_pd, x_pd_, dy_pd, create_graph=True) ddy_pd, = paddle.grad(dx_pd, dy_pd, ddx_pd, create_graph=True) return y_pd, dx_pd, ddy_pd g = paddle.jit.to_static( g, full_graph=True, input_spec=[ InputSpec(shape=[-1, -1, -1], name="x"), InputSpec(shape=[-1, -1, -1, -1], name="mask"), InputSpec(shape=[-1], name="dy_pd"), InputSpec(shape=[-1, -1, -1], name="ddx_pd"), ], ) y_pd, dx_pd, ddy_pd = g(x_pd, mask_pd, dy_pd, ddx_pd) # y_pd, dx_pd = g(x_pd, mask_pd, dy_pd) # 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(), 1e-6, 1e-6) np.testing.assert_allclose(dx_pd.numpy(), dx_pt.detach().cpu().numpy(), 1e-6, 1e-6) np.testing.assert_allclose(ddy_pd.numpy(), ddy_pt.detach().cpu().numpy(), 1e-6, 1e-6)