def test_gather_grad_pynative_int64_fp16():
context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
dim = 0
index = Tensor(np.array([[0, 1, 1, 0, 0], [1, 0, 0, 1, 1]]), ms.int64)
grad = Tensor(
np.array([[0.9031, 0.0890, 0.2779, 0.3198, 0.5710],
[0.6949, 0.8439, 0.2003, 0.6868, 0.4437]]), ms.float16)
expect = np.array([[0.9031, 0.8439, 0.2003, 0.3198, 0.5710],
[0.6949, 0.0890, 0.2779, 0.6868, 0.4437]], np.float16)
output = P.GatherDGrad(dim)(index, grad)
error = 1e-4
diff = output.asnumpy() - expect
assert np.all(diff < error)
def __init__(self, dim=0, shape=None):
super(NetGrad, self).__init__()
self.op = G.GatherDGrad(dim, shape)
def __init__(self, dim=0):
super(GatherDGradNet, self).__init__()
self.gather_d_grad = P.GatherDGrad(dim)
