[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)