[Comp] Support masked_fill_decomp

[HydrogenSulfate]

PR Category

Operator Mechanism

PR Types

New features

Description

Pcard-75624

添加masked_fill的静态图前向组合算子

动转静反向精度验证代码

import paddle
import numpy as np
import torch


paddle.framework.core._set_prim_all_enabled(True)
# paddle.framework.core.set_prim_eager_enabled(True)

x_pd = paddle.randn([300, 400])
x_pd.stop_gradient = False

mask_pd = paddle.randint_like(x_pd, 0, 2).astype("bool")
# mask_pd = paddle.randint(0, 2, [10]).astype("bool")

value_pd = paddle.randn([])
value_pd.stop_gradient = False

dy_pd = paddle.randn_like(x_pd)

def grad_func(x, mask, v, dy):
    y = paddle.masked_fill(x, mask, v)
    return paddle.grad(y, [x, v], dy, create_graph=False)

grad_func = paddle.jit.to_static(grad_func, full_graph=True)
# print(dy_pd)
# print(mask_pd)
dx_pd, dv_pd = grad_func(x_pd, mask_pd, value_pd, dy_pd)
# print(dv_pd)
# exit()
# print(dv_pd.item())

x_pt = torch.from_dlpack(x_pd.detach())
x_pt.requires_grad = True

mask_pt = torch.from_dlpack(mask_pd.detach())

value_pt = torch.from_dlpack(value_pd.detach())
value_pt.requires_grad = True

dy_pt = torch.from_dlpack(dy_pd.detach())
# print(mask_pd.to("float32").sum().item())
# print(dy_pd.sum().item())
# print((dy_pd * mask_pd.astype(dy_pd.dtype)).sum().item())

def torch_masked_fill(x, mask, v):
    y = x.clone()
    y[mask] = v
    return y

y_pt = torch_masked_fill(x_pt, mask_pt, value_pt)

dx_pt, dv_pt = torch.autograd.grad(y_pt, [x_pt, value_pt], dy_pt, create_graph=True)

np.testing.assert_allclose(dx_pd.numpy(), dx_pt.cpu().numpy(), rtol=1e-5, atol=1e-7)
np.testing.assert_allclose(dv_pd.numpy(), dv_pt.cpu().numpy(), rtol=1e-5, atol=1e-7)
print(dv_pd.item(), dv_pt.item())

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