• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020-2021 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
18from mindspore import Tensor
19from mindspore.ops import operations as P
20import mindspore.nn as nn
21import mindspore.context as context
22
23class GatherNdNet(nn.Cell):
24    def __init__(self):
25        super(GatherNdNet, self).__init__()
26        self.gathernd = P.GatherNd()
27
28    def construct(self, x, indices):
29        return self.gathernd(x, indices)
30
31
32def gathernd0(nptype):
33    x = Tensor(np.arange(3 * 2, dtype=nptype).reshape(3, 2))
34    indices = Tensor(np.array([[1, 1], [0, 1]]).astype(np.int32))
35    expect = np.array([3, 1]).astype(nptype)
36
37    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
38    gathernd = GatherNdNet()
39    output = gathernd(x, indices)
40
41    assert np.array_equal(output.asnumpy(), expect)
42
43@pytest.mark.level0
44@pytest.mark.platform_x86_gpu_training
45@pytest.mark.env_onecard
46def test_gathernd0_float64():
47    gathernd0(np.float64)
48
49@pytest.mark.level0
50@pytest.mark.platform_x86_gpu_training
51@pytest.mark.env_onecard
52def test_gathernd0_float32():
53    gathernd0(np.float32)
54
55@pytest.mark.level1
56@pytest.mark.platform_x86_gpu_training
57@pytest.mark.env_onecard
58def test_gathernd0_float16():
59    gathernd0(np.float16)
60
61@pytest.mark.level1
62@pytest.mark.platform_x86_gpu_training
63@pytest.mark.env_onecard
64def test_gathernd0_int32():
65    gathernd0(np.int32)
66
67@pytest.mark.level1
68@pytest.mark.platform_x86_gpu_training
69@pytest.mark.env_onecard
70def test_gathernd0_int16():
71    gathernd0(np.int16)
72
73@pytest.mark.level1
74@pytest.mark.platform_x86_gpu_training
75@pytest.mark.env_onecard
76def test_gathernd0_uint8():
77    gathernd0(np.uint8)
78
79def gathernd1(nptype):
80    x = Tensor(np.arange(2 * 3 * 4 * 5, dtype=nptype).reshape(2, 3, 4, 5))
81    indices = Tensor(np.array([[[[[l, k, j, i] for i in [1, 3, 4]] for j in range(4)]
82                                for k in range(3)] for l in range(2)], dtype='i4'))
83    expect = np.array([[[[1., 3., 4.],
84                         [6., 8., 9.],
85                         [11., 13., 14.],
86                         [16., 18., 19.]],
87
88                        [[21., 23., 24.],
89                         [26., 28., 29.],
90                         [31., 33., 34.],
91                         [36., 38., 39.]],
92
93                        [[41., 43., 44.],
94                         [46., 48., 49.],
95                         [51., 53., 54.],
96                         [56., 58., 59.]]],
97
98                       [[[61., 63., 64.],
99                         [66., 68., 69.],
100                         [71., 73., 74.],
101                         [76., 78., 79.]],
102
103                        [[81., 83., 84.],
104                         [86., 88., 89.],
105                         [91., 93., 94.],
106                         [96., 98., 99.]],
107
108                        [[101., 103., 104.],
109                         [106., 108., 109.],
110                         [111., 113., 114.],
111                         [116., 118., 119.]]]]).astype(nptype)
112
113    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
114    gather = GatherNdNet()
115    output = gather(x, indices)
116
117    assert np.array_equal(output.asnumpy(), expect)
118
119@pytest.mark.level0
120@pytest.mark.platform_x86_gpu_training
121@pytest.mark.env_onecard
122def test_gathernd1_float64():
123    gathernd1(np.float64)
124
125@pytest.mark.level0
126@pytest.mark.platform_x86_gpu_training
127@pytest.mark.env_onecard
128def test_gathernd1_float32():
129    gathernd1(np.float32)
130
131@pytest.mark.level1
132@pytest.mark.platform_x86_gpu_training
133@pytest.mark.env_onecard
134def test_gathernd1_float16():
135    gathernd1(np.float16)
136
137@pytest.mark.level1
138@pytest.mark.platform_x86_gpu_training
139@pytest.mark.env_onecard
140def test_gathernd1_int32():
141    gathernd1(np.int32)
142
143@pytest.mark.level1
144@pytest.mark.platform_x86_gpu_training
145@pytest.mark.env_onecard
146def test_gathernd1_int16():
147    gathernd1(np.int16)
148
149@pytest.mark.level1
150@pytest.mark.platform_x86_gpu_training
151@pytest.mark.env_onecard
152def test_gathernd1_uint8():
153    gathernd1(np.uint8)
154
155def gathernd2(nptype):
156    x = Tensor(np.array([[4., 5., 4., 1., 5.],
157                         [4., 9., 5., 6., 4.],
158                         [9., 8., 4., 3., 6.],
159                         [0., 4., 2., 2., 8.],
160                         [1., 8., 6., 2., 8.],
161                         [8., 1., 9., 7., 3.],
162                         [7., 9., 2., 5., 7.],
163                         [9., 8., 6., 8., 5.],
164                         [3., 7., 2., 7., 4.],
165                         [4., 2., 8., 2., 9.]]).astype(np.float16))
166
167    indices = Tensor(np.array([[4000], [1], [300000]]).astype(np.int32))
168    expect = np.array([[0., 0., 0., 0., 0.],
169                       [4., 9., 5., 6., 4.],
170                       [0., 0., 0., 0., 0.]])
171
172    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
173    gathernd = GatherNdNet()
174    output = gathernd(x, indices)
175
176    assert np.array_equal(output.asnumpy(), expect)
177
178@pytest.mark.level0
179@pytest.mark.platform_x86_gpu_training
180@pytest.mark.env_onecard
181def test_gathernd2_float64():
182    gathernd2(np.float64)
183
184@pytest.mark.level0
185@pytest.mark.platform_x86_gpu_training
186@pytest.mark.env_onecard
187def test_gathernd2_float32():
188    gathernd2(np.float32)
189
190@pytest.mark.level1
191@pytest.mark.platform_x86_gpu_training
192@pytest.mark.env_onecard
193def test_gathernd2_float16():
194    gathernd2(np.float16)
195
196@pytest.mark.level1
197@pytest.mark.platform_x86_gpu_training
198@pytest.mark.env_onecard
199def test_gathernd2_int32():
200    gathernd2(np.int32)
201
202@pytest.mark.level1
203@pytest.mark.platform_x86_gpu_training
204@pytest.mark.env_onecard
205def test_gathernd2_int16():
206    gathernd2(np.int16)
207
208@pytest.mark.level1
209@pytest.mark.platform_x86_gpu_training
210@pytest.mark.env_onecard
211def test_gathernd2_uint8():
212    gathernd2(np.uint8)
213
214@pytest.mark.level1
215@pytest.mark.platform_x86_gpu_training
216@pytest.mark.env_onecard
217def test_gathernd_bool():
218    x = Tensor(np.array([[True, False], [False, False]]).astype(np.bool))
219    indices = Tensor(np.array([[0, 0], [0, 1], [1, 0], [1, 1]]).astype(np.int32))
220    expect = np.array([True, False, False, False]).astype(np.bool)
221
222    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
223    gathernd = GatherNdNet()
224    output = gathernd(x, indices)
225
226    assert np.array_equal(output.asnumpy(), expect)
227
228@pytest.mark.level0
229@pytest.mark.platform_x86_gpu_training
230@pytest.mark.env_onecard
231def test_gathernd_indices_int64():
232    x = Tensor(np.array([[True, False], [False, False]]).astype(np.bool))
233    indices = Tensor(np.array([[0, 0], [0, 1], [1, 0], [1, 1]]).astype(np.int64))
234    expect = np.array([True, False, False, False]).astype(np.bool)
235
236    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
237    gathernd = GatherNdNet()
238    output = gathernd(x, indices)
239
240    assert np.array_equal(output.asnumpy(), expect)
241