• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 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
15import numpy as np
16
17import mindspore as ms
18import mindspore.nn as nn
19from mindspore import Tensor
20from mindspore import context
21from mindspore.common.api import _cell_graph_executor
22from mindspore.ops import composite as C
23from mindspore.ops import operations as P
24from tests.ut.python.ops.test_math_ops import VirtualLoss
25
26
27grad_all = C.GradOperation(get_all=True)
28
29
30class NetWithLoss(nn.Cell):
31    def __init__(self, network):
32        super(NetWithLoss, self).__init__()
33        self.loss = VirtualLoss()
34        self.network = network
35
36    def construct(self, x, y, b):
37        predict = self.network(x, y, b)
38        return self.loss(predict)
39
40
41class GradWrap(nn.Cell):
42    def __init__(self, network):
43        super(GradWrap, self).__init__()
44        self.network = network
45
46    def construct(self, x, y, b):
47        return grad_all(self.network)(x, y, b)
48
49
50def compile_net(net, x, y, b):
51    net.set_auto_parallel()
52    net.set_train()
53    _cell_graph_executor.compile(net, x, y, b)
54
55
56def test_matmul_pow():
57    class Net(nn.Cell):
58        def __init__(self, strategy1, strategy2):
59            super().__init__()
60            self.matmul = P.MatMul().shard(strategy1)
61            self.pow = P.Pow().shard(strategy2)
62            self.matmul2 = P.MatMul().shard(strategy1)
63
64        def construct(self, x, y, b):
65            out = self.matmul(x, y)
66            out = self.pow(out, 2.0)
67            out = self.matmul2(out, b)
68            return out
69
70    context.set_auto_parallel_context(device_num=8, global_rank=0)
71    strategy1 = ((2, 2), (2, 2))
72    strategy2 = ((4, 2), ())
73    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
74    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
75
76    x = Tensor(np.ones([128, 32]), dtype=ms.float32)
77    y = Tensor(np.ones([32, 64]), dtype=ms.float32)
78    b = Tensor(np.ones([64, 64]), dtype=ms.float32)
79    compile_net(net, x, y, b)
80
81
82def test_matmul_exp():
83    class Net(nn.Cell):
84        def __init__(self, strategy1, strategy2):
85            super().__init__()
86            self.matmul = P.MatMul().shard(strategy1)
87            self.exp = P.Exp().shard(strategy2)
88            self.matmul2 = P.MatMul().shard(strategy1)
89
90        def construct(self, x, y, b):
91            out = self.matmul(x, y)
92            out = self.exp(out)
93            out = self.matmul2(out, b)
94            return out
95
96    context.set_auto_parallel_context(device_num=8, global_rank=0)
97    strategy1 = ((2, 2), (2, 2))
98    strategy2 = ((4, 2),)
99    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
100    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
101
102    x = Tensor(np.ones([128, 32]), dtype=ms.float32)
103    y = Tensor(np.ones([32, 64]), dtype=ms.float32)
104    b = Tensor(np.ones([64, 64]), dtype=ms.float32)
105    compile_net(net, x, y, b)
106
107
108def test_matmul_log():
109    class Net(nn.Cell):
110        def __init__(self, strategy1, strategy2):
111            super().__init__()
112            self.matmul = P.MatMul().shard(strategy1)
113            self.log = P.Log().shard(strategy2)
114            self.matmul2 = P.MatMul().shard(strategy1)
115
116        def construct(self, x, y, b):
117            out = self.matmul(x, y)
118            out = self.log(out)
119            out = self.matmul2(out, b)
120            return out
121
122    context.set_auto_parallel_context(device_num=8, global_rank=0)
123    strategy1 = ((2, 2), (2, 2))
124    strategy2 = ((4, 2),)
125    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
126    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
127
128    x = Tensor(np.ones([128, 32]), dtype=ms.float32)
129    y = Tensor(np.ones([32, 64]), dtype=ms.float32)
130    b = Tensor(np.ones([64, 64]), dtype=ms.float32)
131    compile_net(net, x, y, b)
132
133def test_matmul_abs():
134    class Net(nn.Cell):
135        def __init__(self, strategy1, strategy2):
136            super().__init__()
137            self.matmul = P.MatMul().shard(strategy1)
138            self.abs = P.Abs().shard(strategy2)
139            self.matmul2 = P.MatMul().shard(strategy1)
140
141        def construct(self, x, y, b):
142            out = self.matmul(x, y)
143            out = self.abs(out)
144            out = self.matmul2(out, b)
145            return out
146
147    context.set_auto_parallel_context(device_num=8, global_rank=0)
148    strategy1 = ((2, 2), (2, 2))
149    strategy2 = ((4, 2),)
150    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
151    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
152
153    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
154    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
155    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
156    compile_net(net, x, y, b)
157
158def test_matmul_sign():
159    class Net(nn.Cell):
160        def __init__(self, strategy1, strategy2):
161            super().__init__()
162            self.matmul = P.MatMul().shard(strategy1)
163            self.sign = P.Sign().shard(strategy2)
164            self.matmul2 = P.MatMul().shard(strategy1)
165
166        def construct(self, x, y, b):
167            out = self.matmul(x, y)
168            out = self.sign(out)
169            out = self.matmul2(out, b)
170            return out
171
172    context.set_auto_parallel_context(device_num=8, global_rank=0)
173    strategy1 = ((2, 2), (2, 2))
174    strategy2 = ((4, 2),)
175    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
176    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
177
178    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
179    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
180    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
181    compile_net(net, x, y, b)
182
183
184def test_matmul_floor():
185    class Net(nn.Cell):
186        def __init__(self, strategy1, strategy2):
187            super().__init__()
188            self.matmul = P.MatMul().shard(strategy1)
189            self.floor = P.Floor().shard(strategy2)
190            self.matmul2 = P.MatMul().shard(strategy1)
191
192        def construct(self, x, y, b):
193            out = self.matmul(x, y)
194            out = self.floor(out)
195            out = self.matmul2(out, b)
196            return out
197
198    context.set_auto_parallel_context(device_num=8, global_rank=0)
199    strategy1 = ((2, 2), (2, 2))
200    strategy2 = ((4, 2),)
201    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
202    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
203
204    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
205    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
206    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
207    compile_net(net, x, y, b)
208
209def test_matmul_round():
210    class Net(nn.Cell):
211        def __init__(self, strategy1, strategy2):
212            super().__init__()
213            self.matmul = P.MatMul().shard(strategy1)
214            self.round = P.Round().shard(strategy2)
215            self.matmul2 = P.MatMul().shard(strategy1)
216
217        def construct(self, x, y, b):
218            out = self.matmul(x, y)
219            out = self.round(out)
220            out = self.matmul2(out, b)
221            return out
222
223    context.set_auto_parallel_context(device_num=8, global_rank=0)
224    strategy1 = ((2, 2), (2, 2))
225    strategy2 = ((4, 2),)
226    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
227    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
228
229    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
230    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
231    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
232    compile_net(net, x, y, b)
233
234
235def test_matmul_reciprocal():
236    class Net(nn.Cell):
237        def __init__(self, strategy1, strategy2):
238            super().__init__()
239            self.matmul = P.MatMul().shard(strategy1)
240            self.reciprocal = P.Reciprocal().shard(strategy2)
241            self.matmul2 = P.MatMul().shard(strategy1)
242
243        def construct(self, x, y, b):
244            out = self.matmul(x, y)
245            out = self.reciprocal(out)
246            out = self.matmul2(out, b)
247            return out
248
249    context.set_auto_parallel_context(device_num=8, global_rank=0)
250    strategy1 = ((2, 2), (2, 2))
251    strategy2 = ((4, 2),)
252    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
253    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
254
255    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
256    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
257    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
258    compile_net(net, x, y, b)
259
260
261def test_matmul_inv():
262    class Net(nn.Cell):
263        def __init__(self, strategy1, strategy2):
264            super().__init__()
265            self.matmul = P.MatMul().shard(strategy1)
266            self.inv = P.Inv().shard(strategy2)
267            self.matmul2 = P.MatMul().shard(strategy1)
268
269        def construct(self, x, y, b):
270            out = self.matmul(x, y)
271            out = self.inv(out)
272            out = self.matmul2(out, b)
273            return out
274
275    context.set_auto_parallel_context(device_num=8, global_rank=0)
276    strategy1 = ((2, 2), (2, 2))
277    strategy2 = ((4, 2),)
278    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
279    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
280
281    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
282    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
283    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
284    compile_net(net, x, y, b)
285
286
287def test_matmul_rsqrt():
288    class Net(nn.Cell):
289        def __init__(self, strategy1, strategy2):
290            super().__init__()
291            self.matmul = P.MatMul().shard(strategy1)
292            self.rsqrt = P.Rsqrt().shard(strategy2)
293            self.matmul2 = P.MatMul().shard(strategy1)
294
295        def construct(self, x, y, b):
296            out = self.matmul(x, y)
297            out = self.rsqrt(out)
298            out = self.matmul2(out, b)
299            return out
300
301    context.set_auto_parallel_context(device_num=8, global_rank=0)
302    strategy1 = ((2, 2), (2, 2))
303    strategy2 = ((4, 2),)
304    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
305    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
306
307    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
308    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
309    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
310    compile_net(net, x, y, b)
311
312
313def test_matmul_tan():
314    class Net(nn.Cell):
315        def __init__(self, strategy1, strategy2):
316            super().__init__()
317            self.matmul = P.MatMul().shard(strategy1)
318            self.tan = P.Tan().shard(strategy2)
319            self.matmul2 = P.MatMul().shard(strategy1)
320
321        def construct(self, x, y, b):
322            out = self.matmul(x, y)
323            out = self.tan(out)
324            out = self.matmul2(out, b)
325            return out
326
327    context.set_auto_parallel_context(device_num=8, global_rank=0)
328    strategy1 = ((2, 2), (2, 2))
329    strategy2 = ((4, 2),)
330    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
331    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
332
333    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
334    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
335    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
336    compile_net(net, x, y, b)
337
338
339def test_matmul_sin():
340    class Net(nn.Cell):
341        def __init__(self, strategy1, strategy2):
342            super().__init__()
343            self.matmul = P.MatMul().shard(strategy1)
344            self.sin = P.Sin().shard(strategy2)
345            self.matmul2 = P.MatMul().shard(strategy1)
346
347        def construct(self, x, y, b):
348            out = self.matmul(x, y)
349            out = self.sin(out)
350            out = self.matmul2(out, b)
351            return out
352
353    context.set_auto_parallel_context(device_num=8, global_rank=0)
354    strategy1 = ((2, 2), (2, 2))
355    strategy2 = ((4, 2),)
356    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
357    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
358
359    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
360    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
361    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
362    compile_net(net, x, y, b)
363
364
365def test_matmul_sinh():
366    class Net(nn.Cell):
367        def __init__(self, strategy1, strategy2):
368            super().__init__()
369            self.matmul = P.MatMul().shard(strategy1)
370            self.sinh = P.Sinh().shard(strategy2)
371            self.matmul2 = P.MatMul().shard(strategy1)
372
373        def construct(self, x, y, b):
374            out = self.matmul(x, y)
375            out = self.sinh(out)
376            out = self.matmul2(out, b)
377            return out
378
379    context.set_auto_parallel_context(device_num=8, global_rank=0)
380    strategy1 = ((2, 2), (2, 2))
381    strategy2 = ((4, 2),)
382    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
383    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
384
385    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
386    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
387    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
388    compile_net(net, x, y, b)
389
390
391def test_matmul_log1p():
392    class Net(nn.Cell):
393        def __init__(self, strategy1, strategy2):
394            super().__init__()
395            self.matmul = P.MatMul().shard(strategy1)
396            self.log1p = P.Log1p().shard(strategy2)
397            self.matmul2 = P.MatMul().shard(strategy1)
398
399        def construct(self, x, y, b):
400            out = self.matmul(x, y)
401            out = self.log1p(out)
402            out = self.matmul2(out, b)
403            return out
404
405    context.set_auto_parallel_context(device_num=8, global_rank=0)
406    strategy1 = ((2, 2), (2, 2))
407    strategy2 = ((4, 2),)
408    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
409    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
410
411    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
412    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
413    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
414    compile_net(net, x, y, b)
415
416
417def test_matmul_expm1():
418    class Net(nn.Cell):
419        def __init__(self, strategy1, strategy2):
420            super().__init__()
421            self.matmul = P.MatMul().shard(strategy1)
422            self.expm1 = P.Expm1().shard(strategy2)
423            self.matmul2 = P.MatMul().shard(strategy1)
424
425        def construct(self, x, y, b):
426            out = self.matmul(x, y)
427            out = self.expm1(out)
428            out = self.matmul2(out, b)
429            return out
430
431    context.set_auto_parallel_context(device_num=8, global_rank=0)
432    strategy1 = ((2, 2), (2, 2))
433    strategy2 = ((4, 2),)
434    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
435    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
436
437    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
438    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
439    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
440    compile_net(net, x, y, b)
441
442
443def test_matmul_cosh():
444    class Net(nn.Cell):
445        def __init__(self, strategy1, strategy2):
446            super().__init__()
447            self.matmul = P.MatMul().shard(strategy1)
448            self.cosh = P.Cosh().shard(strategy2)
449            self.matmul2 = P.MatMul().shard(strategy1)
450
451        def construct(self, x, y, b):
452            out = self.matmul(x, y)
453            out = self.cosh(out)
454            out = self.matmul2(out, b)
455            return out
456
457    context.set_auto_parallel_context(device_num=8, global_rank=0)
458    strategy1 = ((2, 2), (2, 2))
459    strategy2 = ((4, 2),)
460    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
461    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
462
463    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
464    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
465    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
466    compile_net(net, x, y, b)
467
468def test_matmul_erf():
469    class Net(nn.Cell):
470        def __init__(self, strategy1, strategy2):
471            super().__init__()
472            self.matmul = P.MatMul().shard(strategy1)
473            self.erf = P.Erf().shard(strategy2)
474            self.matmul2 = P.MatMul().shard(strategy1)
475
476        def construct(self, x, y, b):
477            out = self.matmul(x, y)
478            out = self.erf(out)
479            out = self.matmul2(out, b)
480            return out
481
482    context.set_auto_parallel_context(device_num=8, global_rank=0)
483    strategy1 = ((2, 2), (2, 2))
484    strategy2 = ((4, 2),)
485    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
486    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
487
488    x = Tensor(np.random.uniform(1, 5, size=(128, 32)), dtype=ms.float32)
489    y = Tensor(np.random.uniform(1, 5, size=(32, 64)), dtype=ms.float32)
490    b = Tensor(np.random.uniform(1, 5, size=(64, 64)), dtype=ms.float32)
491    compile_net(net, x, y, b)
492
493
494def test_matmul_erfc():
495    class Net(nn.Cell):
496        def __init__(self, strategy1, strategy2):
497            super().__init__()
498            self.matmul = P.MatMul().shard(strategy1)
499            self.erfc = P.Erfc().shard(strategy2)
500            self.matmul2 = P.MatMul().shard(strategy1)
501
502        def construct(self, x, y, b):
503            out = self.matmul(x, y)
504            out = self.erfc(out)
505            out = self.matmul2(out, b)
506            return out
507
508    context.set_auto_parallel_context(device_num=8, global_rank=0)
509    strategy1 = ((2, 2), (2, 2))
510    strategy2 = ((4, 2),)
511    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
512    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
513
514    x = Tensor(np.random.uniform(1, 5, size=(128, 32)), dtype=ms.float32)
515    y = Tensor(np.random.uniform(1, 5, size=(32, 64)), dtype=ms.float32)
516    b = Tensor(np.random.uniform(1, 5, size=(64, 64)), dtype=ms.float32)
517    compile_net(net, x, y, b)
518
519
520def test_matmul_zeroslike():
521    class Net(nn.Cell):
522        def __init__(self, strategy1, strategy2):
523            super().__init__()
524            self.matmul = P.MatMul().shard(strategy1)
525            self.zeroslike = P.ZerosLike().shard(strategy2)
526            self.matmul2 = P.MatMul().shard(strategy1)
527
528        def construct(self, x, y, b):
529            out = self.matmul(x, y)
530            out = self.zeroslike(out)
531            out = self.matmul2(out, b)
532            return out
533
534    context.set_auto_parallel_context(device_num=8, global_rank=0)
535    strategy1 = ((2, 2), (2, 2))
536    strategy2 = ((4, 2),)
537    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
538    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
539
540    x = Tensor(np.random.uniform(1, 5, size=(128, 32)), dtype=ms.float32)
541    y = Tensor(np.random.uniform(1, 5, size=(32, 64)), dtype=ms.float32)
542    b = Tensor(np.random.uniform(1, 5, size=(64, 64)), dtype=ms.float32)
543    compile_net(net, x, y, b)
544
545
546def test_matmul_oneslike():
547    class Net(nn.Cell):
548        def __init__(self, strategy1, strategy2):
549            super().__init__()
550            self.matmul = P.MatMul().shard(strategy1)
551            self.oneslike = P.OnesLike().shard(strategy2)
552            self.matmul2 = P.MatMul().shard(strategy1)
553
554        def construct(self, x, y, b):
555            out = self.matmul(x, y)
556            out = self.oneslike(out)
557            out = self.matmul2(out, b)
558            return out
559
560    context.set_auto_parallel_context(device_num=8, global_rank=0)
561    strategy1 = ((2, 2), (2, 2))
562    strategy2 = ((4, 2),)
563    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
564    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
565
566    x = Tensor(np.random.uniform(1, 5, size=(128, 32)), dtype=ms.float32)
567    y = Tensor(np.random.uniform(1, 5, size=(32, 64)), dtype=ms.float32)
568    b = Tensor(np.random.uniform(1, 5, size=(64, 64)), dtype=ms.float32)
569    compile_net(net, x, y, b)
570
571
572def test_matmul_BesselI0e():
573    class Net(nn.Cell):
574        def __init__(self, strategy1, strategy2):
575            super().__init__()
576            self.matmul = P.MatMul().shard(strategy1)
577            self.BesselI0e = P.BesselI0e().shard(strategy2)
578            self.matmul2 = P.MatMul().shard(strategy1)
579
580        def construct(self, x, y, b):
581            out = self.matmul(x, y)
582            out = self.BesselI0e(out)
583            out = self.matmul2(out, b)
584            return out
585
586    context.set_auto_parallel_context(device_num=8, global_rank=0)
587    strategy1 = ((2, 2), (2, 2))
588    strategy2 = ((4, 2),)
589    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
590    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
591
592    x = Tensor(np.random.uniform(1, 5, size=(128, 32)), dtype=ms.float32)
593    y = Tensor(np.random.uniform(1, 5, size=(32, 64)), dtype=ms.float32)
594    b = Tensor(np.random.uniform(1, 5, size=(64, 64)), dtype=ms.float32)
595    compile_net(net, x, y, b)
596
597
598def test_matmul_BesselI1e():
599    class Net(nn.Cell):
600        def __init__(self, strategy1, strategy2):
601            super().__init__()
602            self.matmul = P.MatMul().shard(strategy1)
603            self.BesselI1e = P.BesselI1e().shard(strategy2)
604            self.matmul2 = P.MatMul().shard(strategy1)
605
606        def construct(self, x, y, b):
607            out = self.matmul(x, y)
608            out = self.BesselI1e(out)
609            out = self.matmul2(out, b)
610            return out
611
612    context.set_auto_parallel_context(device_num=8, global_rank=0)
613    strategy1 = ((2, 2), (2, 2))
614    strategy2 = ((4, 2),)
615    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
616    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
617
618    x = Tensor(np.random.uniform(1, 5, size=(128, 32)), dtype=ms.float32)
619    y = Tensor(np.random.uniform(1, 5, size=(32, 64)), dtype=ms.float32)
620    b = Tensor(np.random.uniform(1, 5, size=(64, 64)), dtype=ms.float32)
621    compile_net(net, x, y, b)
622
623
624def test_matmul_ceil():
625    class Net(nn.Cell):
626        def __init__(self, strategy1, strategy2):
627            super().__init__()
628            self.matmul = P.MatMul().shard(strategy1)
629            self.Ceil = P.Ceil().shard(strategy2)
630            self.matmul2 = P.MatMul().shard(strategy1)
631
632        def construct(self, x, y, b):
633            out = self.matmul(x, y)
634            out = self.Ceil(out)
635            out = self.matmul2(out, b)
636            return out
637
638    context.set_auto_parallel_context(device_num=8, global_rank=0)
639    strategy1 = ((2, 2), (2, 2))
640    strategy2 = ((4, 2),)
641    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
642    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
643
644    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
645    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
646    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
647    compile_net(net, x, y, b)
648
649
650def test_matmul_atan():
651    class Net(nn.Cell):
652        def __init__(self, strategy1, strategy2):
653            super().__init__()
654            self.matmul = P.MatMul().shard(strategy1)
655            self.atan = P.Atan().shard(strategy2)
656            self.matmul2 = P.MatMul().shard(strategy1)
657
658        def construct(self, x, y, b):
659            out = self.matmul(x, y)
660            out = self.atan(out)
661            out = self.matmul2(out, b)
662            return out
663
664    context.set_auto_parallel_context(device_num=8, global_rank=0)
665    strategy1 = ((2, 2), (2, 2))
666    strategy2 = ((4, 2),)
667    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
668    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
669
670    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
671    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
672    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
673    compile_net(net, x, y, b)
674
675
676def test_matmul_Atanh():
677    class Net(nn.Cell):
678        def __init__(self, strategy1, strategy2):
679            super().__init__()
680            self.matmul = P.MatMul().shard(strategy1)
681            self.atanh = P.Atanh().shard(strategy2)
682            self.matmul2 = P.MatMul().shard(strategy1)
683
684        def construct(self, x, y, b):
685            out = self.matmul(x, y)
686            out = self.atanh(out)
687            out = self.matmul2(out, b)
688            return out
689
690    context.set_auto_parallel_context(device_num=8, global_rank=0)
691    strategy1 = ((2, 2), (2, 2))
692    strategy2 = ((4, 2),)
693    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
694    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
695
696    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
697    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
698    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
699    compile_net(net, x, y, b)
700
701
702def test_matmul_asin():
703    class Net(nn.Cell):
704        def __init__(self, strategy1, strategy2):
705            super().__init__()
706            self.matmul = P.MatMul().shard(strategy1)
707            self.asin = P.Asin().shard(strategy2)
708            self.matmul2 = P.MatMul().shard(strategy1)
709
710        def construct(self, x, y, b):
711            out = self.matmul(x, y)
712            out = self.asin(out)
713            out = self.matmul2(out, b)
714            return out
715
716    context.set_auto_parallel_context(device_num=8, global_rank=0)
717    strategy1 = ((2, 2), (2, 2))
718    strategy2 = ((4, 2),)
719    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
720    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
721
722    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
723    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
724    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
725    compile_net(net, x, y, b)
726
727
728def test_matmul_asinh():
729    class Net(nn.Cell):
730        def __init__(self, strategy1, strategy2):
731            super().__init__()
732            self.matmul = P.MatMul().shard(strategy1)
733            self.asinh = P.Asinh().shard(strategy2)
734            self.matmul2 = P.MatMul().shard(strategy1)
735
736        def construct(self, x, y, b):
737            out = self.matmul(x, y)
738            out = self.asinh(out)
739            out = self.matmul2(out, b)
740            return out
741
742    context.set_auto_parallel_context(device_num=8, global_rank=0)
743    strategy1 = ((2, 2), (2, 2))
744    strategy2 = ((4, 2),)
745    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
746    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
747
748    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
749    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
750    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
751    compile_net(net, x, y, b)
752
753
754def test_matmul_acosh():
755    class Net(nn.Cell):
756        def __init__(self, strategy1, strategy2):
757            super().__init__()
758            self.matmul = P.MatMul().shard(strategy1)
759            self.acosh = P.Acosh().shard(strategy2)
760            self.matmul2 = P.MatMul().shard(strategy1)
761
762        def construct(self, x, y, b):
763            out = self.matmul(x, y)
764            out = self.acosh(out)
765            out = self.matmul2(out, b)
766            return out
767
768    context.set_auto_parallel_context(device_num=8, global_rank=0)
769    strategy1 = ((2, 2), (2, 2))
770    strategy2 = ((4, 2),)
771    net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))
772    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
773
774    x = Tensor(np.random.uniform(-5, 5, size=(128, 32)), dtype=ms.float32)
775    y = Tensor(np.random.uniform(-5, 5, size=(32, 64)), dtype=ms.float32)
776    b = Tensor(np.random.uniform(-5, 5, size=(64, 64)), dtype=ms.float32)
777    compile_net(net, x, y, b)
778
779
780def test_matmul_logical_not():
781    class Net(nn.Cell):
782        def __init__(self, strategy1, strategy2, strategy3):
783            super().__init__()
784            self.matmul = P.MatMul().shard(strategy1)
785            self.logicalnot = P.LogicalNot().shard(strategy2)
786            self.equal = P.Equal().shard(strategy3)
787
788        def construct(self, x, y, b):
789            out = self.matmul(x, y)
790            out = self.equal(out, b)
791            out = self.logicalnot(out)
792            return out
793
794    context.set_auto_parallel_context(device_num=8, global_rank=0)
795    strategy1 = ((2, 2), (2, 2))
796    strategy2 = ((4, 2),)
797    strategy3 = ((4, 2), (4, 2))
798    net = GradWrap(NetWithLoss(Net(strategy1, strategy2, strategy3)))
799    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
800
801    x = Tensor(np.ones([128, 32]), dtype=ms.float32)
802    y = Tensor(np.ones([32, 64]), dtype=ms.float32)
803    b = Tensor(np.ones([128, 64]), dtype=ms.float32)
804    compile_net(net, x, y, b)
805
806
807def test_matmul_cast():
808    class Net(nn.Cell):
809        def __init__(self, strategy1, strategy2, strategy3):
810            super().__init__()
811            self.matmul = P.MatMul().shard(strategy1)
812            self.cast = P.Cast().shard(strategy2)
813            self.matmul2 = P.MatMul().shard(strategy3)
814
815        def construct(self, x, y, b):
816            out = self.matmul(x, y)
817            b = self.cast(b, ms.float32)
818            out = self.matmul2(out, b)
819            return out
820
821    context.set_auto_parallel_context(device_num=8, global_rank=0)
822    strategy1 = ((2, 2), (2, 2))
823    strategy2 = ((4, 2),)
824    strategy3 = ((1, 4), (4, 2))
825    net = GradWrap(NetWithLoss(Net(strategy1, strategy2, strategy3)))
826    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
827
828    x = Tensor(np.ones([128, 32]), dtype=ms.float32)
829    y = Tensor(np.ones([32, 64]), dtype=ms.float32)
830    b = Tensor(np.ones([64, 64]), dtype=ms.int32)
831    compile_net(net, x, y, b)
832
833
834def test_gradient_fp32_sync():
835    class Net(nn.Cell):
836        def __init__(self, strategy1):
837            super().__init__()
838            self.matmul = P.MatMul().shard(strategy1)
839            self.cast = P.Cast()
840
841        def construct(self, x, y, b):
842            out = self.matmul(x, y)
843            b = self.cast(b, ms.float32)
844            out = self.matmul(out, b)
845            return out
846
847    context.set_auto_parallel_context(device_num=8, global_rank=0, gradient_fp32_sync=True)
848    strategy1 = ((2, 2), (2, 2))
849    net = GradWrap(NetWithLoss(Net(strategy1)))
850    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
851
852    x = Tensor(np.ones([128, 32]), dtype=ms.float32)
853    y = Tensor(np.ones([32, 64]), dtype=ms.float32)
854    b = Tensor(np.ones([64, 64]), dtype=ms.float16)
855    compile_net(net, x, y, b)
856
857
858def test_gradient_fp32_sync1():
859    class Net(nn.Cell):
860        def __init__(self, strategy1):
861            super().__init__()
862            self.matmul = P.MatMul().shard(strategy1)
863            self.cast = P.Cast()
864
865        def construct(self, x, y, b):
866            out = self.matmul(x, y)
867            b = self.cast(b, ms.float16)
868            out = self.matmul(out, b)
869            return out
870
871    context.set_auto_parallel_context(device_num=8, global_rank=0, gradient_fp32_sync=True)
872    strategy1 = ((2, 2), (2, 2))
873    net = GradWrap(NetWithLoss(Net(strategy1)))
874    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
875
876    x = Tensor(np.ones([128, 32]), dtype=ms.float16)
877    y = Tensor(np.ones([32, 64]), dtype=ms.float16)
878    b = Tensor(np.ones([64, 64]), dtype=ms.float32)
879    compile_net(net, x, y, b)
880
881
882def test_gradient_fp32_sync2():
883    class Net(nn.Cell):
884        def __init__(self, strategy1):
885            super().__init__()
886            self.matmul = P.MatMul().shard(strategy1)
887            self.cast = P.Cast()
888
889        def construct(self, x, y, b):
890            out = self.matmul(x, y)
891            b = self.cast(b, ms.float16)
892            out = self.matmul(out, b)
893            return out
894
895    context.set_auto_parallel_context(device_num=8, global_rank=0, gradient_fp32_sync=False)
896    strategy1 = ((2, 2), (2, 2))
897    net = GradWrap(NetWithLoss(Net(strategy1)))
898    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
899
900    x = Tensor(np.ones([128, 32]), dtype=ms.float16)
901    y = Tensor(np.ones([32, 64]), dtype=ms.float16)
902    b = Tensor(np.ones([64, 64]), dtype=ms.float32)
903    compile_net(net, x, y, b)
904
905
906def test_gradient_fp32_sync3():
907    class Net(nn.Cell):
908        def __init__(self, strategy1):
909            super().__init__()
910            self.matmul = P.MatMul().shard(strategy1)
911            self.cast = P.Cast()
912
913        def construct(self, x, y, b):
914            out = self.matmul(x, y)
915            b = self.cast(b, ms.float16)
916            out = self.matmul(out, b)
917            return out
918
919    context.set_auto_parallel_context(device_num=8, global_rank=0)
920    strategy1 = ((2, 2), (2, 2))
921    net = GradWrap(NetWithLoss(Net(strategy1)))
922    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
923
924    x = Tensor(np.ones([128, 32]), dtype=ms.float16)
925    y = Tensor(np.ones([32, 64]), dtype=ms.float16)
926    b = Tensor(np.ones([64, 64]), dtype=ms.float32)
927    compile_net(net, x, y, b)
928
929
930def test_mul_two_cast():
931    class Net(nn.Cell):
932        def __init__(self, strategy1, strategy2, strategy3):
933            super().__init__()
934            self.mul = P.Mul().shard(strategy1)
935            self.mul2 = P.Mul().shard(strategy2)
936            self.cast = P.Cast().shard(strategy3)
937            self.cast2 = P.Cast().shard(strategy3)
938
939        def construct(self, x, y, b):
940            out = self.mul(x, y)
941            out = self.mul2(out, b)
942            out = self.cast(out, ms.int32)
943            out = self.cast2(out, ms.bool_)
944            return out
945
946    context.set_auto_parallel_context(device_num=8, global_rank=0)
947    strategy1 = ((2, 2), (2, 2))
948    strategy2 = ((8, 1), (8, 1))
949    strategy3 = ((8, 1),)
950    net = GradWrap(Net(strategy1, strategy2, strategy3))
951    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
952
953    x = Tensor(np.ones([128, 32]), dtype=ms.float32)
954    y = Tensor(np.ones([128, 32]), dtype=ms.float32)
955    b = Tensor(np.ones([128, 32]), dtype=ms.float32)
956    compile_net(net, x, y, b)
957