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)