[Prim] Add gather nd double grad

[HydrogenSulfate]

PR Category

Operator Mechanism

PR Types

Improvements

Description

Pcard-75624

1. 添加动态图二阶反向组合gather_nd_double_grad，支持非组合模式下gather_nd任意高阶微分。精度测试脚本如下
2. index_put_double_grad的参数名字按照已有算子的写法，从grad_out改为更合理的out_grad

def test_gather_nd_backward():
    import paddle

    for x_shape, index_shape in [
        # ([16], [10]),
        # ([16, 16], [20, 2]),
        # ([12, 13, 14], [88, 1]),
        # ([12, 13, 14], [88, 2]),
        # ([12, 13, 14], [88, 3]),
        # ([128, 32, 32], [128 * 32 * 32]),
        # ([128, 32, 32], [128 * 32 * 32, 1]),
        # ([512, 32, 32], [512 * 32 * 32, 2]),
        ([512, 32, 32], [512 * 32 * 32, 3]),
    ]:
        x_np = np.random.randn(*x_shape).astype("float32")
        n_index = index_shape[0]
        index_dim_size = index_shape[1] if len(index_shape) > 1 else 1
        index_np = [
            np.random.randint(-x_shape[i], x_shape[i], size=[n_index]) # [D]
            for i in range(index_dim_size)
        ]
        index_np = np.stack(index_np, axis=-1)  # [N, D]

        x = paddle.to_tensor(x_np, stop_gradient=False)
        index = paddle.to_tensor(index_np)

        core.set_prim_eager_enabled(True)
        core._set_prim_all_enabled(True)

        # forward
        out = paddle.gather_nd(x, index)
        dout_np = np.random.randn(*out.shape).astype("float32")
        dout = paddle.to_tensor(dout_np, stop_gradient=False)

        # 1-order backward
        dx = paddle.grad(out, x, dout, create_graph=True)[0]
        ddx_np = np.random.randn(*x_shape)
        ddx = paddle.to_tensor(ddx_np, stop_gradient=False)

        # 2-order backward
        ddout = paddle.grad(dx, dout, ddx, create_graph=True)[0]

        # add gather_nd_grad to blacklist to allow execution of gather_nd_double_grad
        core._set_prim_backward_blacklist("gather_nd_grad")

        # forward
        out_ = paddle.gather_nd(x, index)
        # 1-order backward
        dx_ = paddle.grad(out_, x, dout, create_graph=True)[0]

        # 2-order backward
        ddout_ = paddle.grad(dx_, dout, ddx, create_graph=True)[0]

        # compare result
        np.testing.assert_allclose(dx.numpy(), dx_.numpy(), 1e-5, 1e-5)
        np.testing.assert_allclose(ddout.numpy(), ddout_.numpy(), 1e-5, 1e-5)

if __name__ == "__main__":
    test_gather_nd_backward()

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