1# Copyright 2019-2020 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================ 15 16import numpy as np 17import pytest 18 19import mindspore.context as context 20import mindspore.nn as nn 21import mindspore as ms 22from mindspore import Tensor 23from mindspore.ops import operations as P 24 25 26class GatherNet(nn.Cell): 27 def __init__(self, dim=0): 28 super(GatherNet, self).__init__() 29 self.gather = P.GatherD() 30 self.dim = dim 31 32 def construct(self, x, index): 33 return self.gather(x, self.dim, index) 34 35@pytest.mark.level0 36@pytest.mark.platform_x86_gpu_training 37@pytest.mark.env_onecard 38def test_gather_pynative_fp32_int32(): 39 context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU") 40 error = 1e-3 41 x = Tensor(np.array([[1.303, 2.333], [3.232, 4.235]]), ms.float32) 42 dim = 1 43 index = Tensor(np.array([[0, 0], [1, 0]]), ms.int32) 44 expect = np.array([[1.303, 1.303], [4.235, 3.232]], np.float32) 45 output = P.GatherD()(x, dim, index) 46 diff = output.asnumpy() - expect 47 assert np.all(diff < error) 48 49@pytest.mark.level0 50@pytest.mark.platform_x86_gpu_training 51@pytest.mark.env_onecard 52def test_gather_pynative_fp32_int64(): 53 context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU") 54 error = 1e-3 55 x = Tensor(np.array([[1.303, 2.333], [3.232, 4.235]]), ms.float32) 56 dim = 1 57 index = Tensor(np.array([[0, 0], [1, 0]]), ms.int64) 58 expect = np.array([[1.303, 1.303], [4.235, 3.232]], np.float32) 59 output = P.GatherD()(x, dim, index) 60 diff = output.asnumpy() - expect 61 assert np.all(diff < error) 62 63@pytest.mark.level0 64@pytest.mark.platform_x86_gpu_training 65@pytest.mark.env_onecard 66def test_gather_pynative_fp16_int32(): 67 context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU") 68 error = 1e-3 69 x = Tensor(np.array([[1.303, 2.333], [3.232, 4.235]]), ms.float16) 70 dim = 1 71 index = Tensor(np.array([[0, 0], [1, 0]]), ms.int32) 72 expect = np.array([[1.303, 1.303], [4.235, 3.232]], np.float16) 73 output = P.GatherD()(x, dim, index) 74 diff = output.asnumpy() - expect 75 assert np.all(diff < error) 76 77@pytest.mark.level0 78@pytest.mark.platform_x86_gpu_training 79@pytest.mark.env_onecard 80def test_gather_pynative_fp16_int64(): 81 context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU") 82 error = 1e-3 83 x = Tensor(np.array([[1.303, 2.333], [3.232, 4.235]]), ms.float16) 84 dim = 1 85 index = Tensor(np.array([[0, 0], [1, 0]]), ms.int64) 86 expect = np.array([[1.303, 1.303], [4.235, 3.232]], np.float16) 87 output = P.GatherD()(x, dim, index) 88 diff = output.asnumpy() - expect 89 assert np.all(diff < error) 90 91def test_gather_graph_fp32_int32(): 92 context.set_context(mode=context.GRAPH_MODE, device_target="GPU") 93 error = 1e-3 94 x = Tensor(np.array([[1.303, 2.333], [3.232, 4.235]]), ms.float32) 95 dim = 1 96 index = Tensor(np.array([[0, 0], [1, 0]]), ms.int32) 97 expect = np.array([[1.303, 1.303], [4.235, 3.232]], np.float32) 98 gather = GatherNet(dim) 99 output = gather(x, index) 100 diff = output.asnumpy() - expect 101 assert np.all(diff < error) 102 103def test_gather_graph_fp32_int64(): 104 context.set_context(mode=context.GRAPH_MODE, device_target="GPU") 105 error = 1e-3 106 x = Tensor(np.array([[1.303, 2.333], [3.232, 4.235]]), ms.float32) 107 dim = 1 108 index = Tensor(np.array([[0, 0], [1, 0]]), ms.int64) 109 expect = np.array([[1.303, 1.303], [4.235, 3.232]], np.float32) 110 gather = GatherNet(dim) 111 output = gather(x, index) 112 diff = output.asnumpy() - expect 113 assert np.all(diff < error) 114 115def test_gather_graph_fp16_int32(): 116 context.set_context(mode=context.GRAPH_MODE, device_target="GPU") 117 error = 1e-3 118 x = Tensor(np.array([[1.303, 2.333], [3.232, 4.235]]), ms.float16) 119 dim = 1 120 index = Tensor(np.array([[0, 0], [1, 0]]), ms.int32) 121 expect = np.array([[1.303, 1.303], [4.235, 3.232]], np.float16) 122 gather = GatherNet(dim) 123 output = gather(x, index) 124 diff = output.asnumpy() - expect 125 assert np.all(diff < error) 126 127def test_gather_graph_fp16_int64(): 128 context.set_context(mode=context.GRAPH_MODE, device_target="GPU") 129 error = 1e-3 130 x = Tensor(np.array([[1.303, 2.333], [3.232, 4.235]]), ms.float16) 131 dim = 1 132 index = Tensor(np.array([[0, 0], [1, 0]]), ms.int64) 133 expect = np.array([[1.303, 1.303], [4.235, 3.232]], np.float16) 134 gather = GatherNet(dim) 135 output = gather(x, index) 136 diff = output.asnumpy() - expect 137 assert np.all(diff < error) 138