zzy34407230
/
mindspore2022
Not watched
1
0
Fork 0
Code
Releases 22 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: cdd3e50e98
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'cdd3e50e98'
${ noResults }
mindspore2022/tests/ut/python/parallel/test_arithmetic.py

1199 lines
42 kB
Raw Blame History 

  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. 

  15. import numpy as np

  16. 

  17. import mindspore as ms

  18. import mindspore.nn as nn

  19. from mindspore import Parameter, Tensor, context

  20. from mindspore.common.api import _cell_graph_executor

  21. from mindspore.ops import composite as C

  22. from mindspore.ops import operations as P

  23. from tests.ut.python.ops.test_math_ops import VirtualLoss

  24. 

  25. 

  26. grad_all = C.GradOperation(get_all=True)

  27. 

  28. 

  29. class NetWithLoss(nn.Cell):

  30.     def __init__(self, network):

  31.         super(NetWithLoss, self).__init__()

  32.         self.loss = VirtualLoss()

  33.         self.network = network

  34. 

  35.     def construct(self, x, y, b):

  36.         predict = self.network(x, y, b)

  37.         return self.loss(predict)

  38. 

  39. 

  40. class GradWrap(nn.Cell):

  41.     def __init__(self, network):

  42.         super(GradWrap, self).__init__()

  43.         self.network = network

  44. 

  45.     def construct(self, x, y, b):

  46.         return grad_all(self.network)(x, y, b)

  47. 

  48. 

  49. def compile_net(net, x, y, b):

  50.     net.set_auto_parallel()

  51.     net.set_train()

  52.     _cell_graph_executor.compile(net, x, y, b)

  53. 

  54. 

  55. def test_matmul_sub():

  56.     """

  57.     Feature: distribute operator sub in auto parallel.

  58.     Description: matmul-sub net with strategy in semi auto parallel.

  59.     Expectation: compile done without error.

  60.     """

  61.     class Net(nn.Cell):

  62.         def __init__(self, strategy1, strategy2):

  63.             super().__init__()

  64.             self.matmul = P.MatMul().shard(strategy1)

  65.             self.sub = P.Sub().shard(strategy2)

  66. 

  67.         def construct(self, x, y, b):

  68.             out = self.matmul(x, y)

  69.             out = self.sub(out, b)

  70.             return out

  71. 

  72.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  73.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  74.     strategy1 = ((2, 2), (2, 2))

  75.     strategy2 = ((4, 2), (4, 2))

  76.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  77. 

  78.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  79.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  80.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  81.     compile_net(net, x, y, b)

  82. 

  83. 

  84. def test_matmul_add():

  85.     """

  86.     Feature: distribute operator sub in auto parallel.

  87.     Description: matmul-add net with strategy in semi auto parallel.

  88.     Expectation: compile done without error.

  89.     """

  90.     class Net(nn.Cell):

  91.         def __init__(self, strategy1, strategy2):

  92.             super().__init__()

  93.             self.matmul = P.MatMul().shard(strategy1)

  94.             self.add = P.Add().shard(strategy2)

  95. 

  96.         def construct(self, x, y, b):

  97.             out = self.matmul(x, y)

  98.             out = self.add(out, b)

  99.             return out

  100. 

  101.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  102.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  103.     strategy1 = ((2, 2), (2, 2))

  104.     strategy2 = ((4, 2), (4, 2))

  105.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  106. 

  107.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  108.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  109.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  110.     compile_net(net, x, y, b)

  111. 

  112. 

  113. def test_matmul_mul():

  114.     """

  115.     Feature: distribute operator sub in auto parallel.

  116.     Description: matmul-mul net with strategy in semi auto parallel.

  117.     Expectation: compile done without error.

  118.     """

  119.     class Net(nn.Cell):

  120.         def __init__(self, strategy1, strategy2):

  121.             super().__init__()

  122.             self.matmul = P.MatMul().shard(strategy1)

  123.             self.mul = P.Mul().shard(strategy2)

  124. 

  125.         def construct(self, x, y, b):

  126.             out = self.matmul(x, y)

  127.             out = self.mul(out, b)

  128.             return out

  129. 

  130.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  131.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  132.     strategy1 = ((2, 2), (2, 2))

  133.     strategy2 = ((4, 2), (4, 2))

  134.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  135. 

  136.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  137.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  138.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  139.     compile_net(net, x, y, b)

  140. 

  141. def test_matmul_mod():

  142.     """

  143.     Feature: distribute operator sub in auto parallel.

  144.     Description: matmul-mod net with strategy in semi auto parallel.

  145.     Expectation: compile done without error.

  146.     """

  147.     class Net(nn.Cell):

  148.         def __init__(self, strategy1, strategy2):

  149.             super().__init__()

  150.             self.matmul = P.MatMul().shard(strategy1)

  151.             self.mod = P.Mod().shard(strategy2)

  152. 

  153.         def construct(self, x, y, b):

  154.             out = self.matmul(x, y)

  155.             out = self.mod(out, b)

  156.             return out

  157. 

  158.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  159.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  160.     strategy1 = ((2, 2), (2, 2))

  161.     strategy2 = ((4, 2), (4, 2))

  162.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  163. 

  164.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  165.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  166.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  167.     compile_net(net, x, y, b)

  168. 

  169. def test_matmul_floormod():

  170.     """

  171.     Feature: distribute operator sub in auto parallel.

  172.     Description: matmul-floormod net with strategy in semi auto parallel.

  173.     Expectation: compile done without error.

  174.     """

  175.     class Net(nn.Cell):

  176.         def __init__(self, strategy1, strategy2):

  177.             super().__init__()

  178.             self.matmul = P.MatMul().shard(strategy1)

  179.             self.floormod = P.FloorMod().shard(strategy2)

  180. 

  181.         def construct(self, x, y, b):

  182.             out = self.matmul(x, y)

  183.             out = self.floormod(out, b)

  184.             return out

  185. 

  186.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  187.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  188.     strategy1 = ((2, 2), (2, 2))

  189.     strategy2 = ((4, 2), (4, 2))

  190.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  191. 

  192.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  193.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  194.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  195.     compile_net(net, x, y, b)

  196. 

  197. 

  198. def test_matmul_atan2():

  199.     """

  200.     Feature: distribute operator sub in auto parallel.

  201.     Description: matmul-atan2 net with strategy in semi auto parallel.

  202.     Expectation: compile done without error.

  203.     """

  204.     class Net(nn.Cell):

  205.         def __init__(self, strategy1, strategy2):

  206.             super().__init__()

  207.             self.matmul = P.MatMul().shard(strategy1)

  208.             self.atan2 = P.Atan2().shard(strategy2)

  209. 

  210.         def construct(self, x, y, b):

  211.             out = self.matmul(x, y)

  212.             out = self.atan2(out, b)

  213.             return out

  214. 

  215.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  216.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  217.     strategy1 = ((2, 2), (2, 2))

  218.     strategy2 = ((4, 2), (4, 2))

  219.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  220. 

  221.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  222.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  223.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  224.     compile_net(net, x, y, b)

  225. 

  226. 

  227. def test_matmul_divNoNan():

  228.     """

  229.     Feature: distribute operator sub in auto parallel.

  230.     Description: matmul-divNoNan net with strategy in semi auto parallel.

  231.     Expectation: compile done without error.

  232.     """

  233.     class Net(nn.Cell):

  234.         def __init__(self, strategy1, strategy2):

  235.             super().__init__()

  236.             self.matmul = P.MatMul().shard(strategy1)

  237.             self.divNoNan = P.DivNoNan().shard(strategy2)

  238. 

  239.         def construct(self, x, y, b):

  240.             out = self.matmul(x, y)

  241.             out = self.divNoNan(out, b)

  242.             return out

  243. 

  244.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  245.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  246.     strategy1 = ((2, 2), (2, 2))

  247.     strategy2 = ((4, 2), (4, 2))

  248.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  249. 

  250.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  251.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  252.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  253.     compile_net(net, x, y, b)

  254. 

  255. 

  256. def test_matmul_logicaland():

  257.     """

  258.     Feature: distribute operator sub in auto parallel.

  259.     Description: matmul-logical_and net with strategy in semi auto parallel.

  260.     Expectation: compile done without error.

  261.     """

  262.     class Net(nn.Cell):

  263.         def __init__(self, strategy1, strategy2):

  264.             super().__init__()

  265.             self.matmul = P.MatMul().shard(strategy1)

  266.             self.equal = P.Equal().shard(strategy2)

  267.             self.notequal = P.NotEqual().shard(strategy2)

  268.             self.logical = P.LogicalAnd().shard(strategy2)

  269. 

  270.         def construct(self, x, y, b):

  271.             out = self.matmul(x, y)

  272.             out1 = self.equal(out, b)

  273.             out = self.matmul(x, y)

  274.             out2 = self.notequal(out, b)

  275.             out = self.logical(out1, out2)

  276.             return out

  277. 

  278.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  279.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  280.     strategy1 = ((2, 2), (2, 2))

  281.     strategy2 = ((4, 2), (4, 2))

  282.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  283. 

  284.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  285.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  286.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  287.     compile_net(net, x, y, b)

  288. 

  289. 

  290. def test_matmul_logicalor():

  291.     """

  292.     Feature: distribute operator sub in auto parallel.

  293.     Description: matmul-logical_or net with strategy in semi auto parallel.

  294.     Expectation: compile done without error.

  295.     """

  296.     class Net(nn.Cell):

  297.         def __init__(self, strategy1, strategy2):

  298.             super().__init__()

  299.             self.matmul = P.MatMul().shard(strategy1)

  300.             self.equal = P.Equal().shard(strategy2)

  301.             self.notequal = P.NotEqual().shard(strategy2)

  302.             self.logical = P.LogicalOr().shard(strategy2)

  303. 

  304.         def construct(self, x, y, b):

  305.             out = self.matmul(x, y)

  306.             out1 = self.equal(out, b)

  307.             out = self.matmul(x, y)

  308.             out2 = self.notequal(out, b)

  309.             out = self.logical(out1, out2)

  310.             return out

  311. 

  312.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  313.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  314.     strategy1 = ((2, 2), (2, 2))

  315.     strategy2 = ((4, 2), (4, 2))

  316.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  317. 

  318.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  319.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  320.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  321.     compile_net(net, x, y, b)

  322. 

  323. 

  324. def test_matmul_div():

  325.     """

  326.     Feature: distribute operator sub in auto parallel.

  327.     Description: matmul-div net with strategy in semi auto parallel.

  328.     Expectation: compile done without error.

  329.     """

  330.     class Net(nn.Cell):

  331.         def __init__(self, strategy1, strategy2):

  332.             super().__init__()

  333.             self.matmul = P.MatMul().shard(strategy1)

  334.             self.div = P.Div().shard(strategy2)

  335. 

  336.         def construct(self, x, y, b):

  337.             out = self.matmul(x, y)

  338.             out = self.div(out, b)

  339.             return out

  340. 

  341.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  342.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  343.     strategy1 = ((2, 2), (2, 2))

  344.     strategy2 = ((4, 2), (4, 2))

  345.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  346. 

  347.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  348.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  349.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  350.     compile_net(net, x, y, b)

  351. 

  352. 

  353. def test_matmul_add_broadcast():

  354.     """

  355.     Feature: distribute operator sub in auto parallel.

  356.     Description: matmul-add broadcast net with strategy in semi auto parallel.

  357.     Expectation: compile done without error.

  358.     """

  359.     class Net(nn.Cell):

  360.         def __init__(self, strategy1, strategy2):

  361.             super().__init__()

  362.             self.matmul = P.MatMul().shard(strategy1)

  363.             self.add = P.Add().shard(strategy2)

  364. 

  365.         def construct(self, x, y, b):

  366.             out = self.matmul(x, y)

  367.             out = self.add(out, b)

  368.             return out

  369. 

  370.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  371.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  372.     strategy1 = ((2, 2), (2, 2))

  373.     strategy2 = ((4, 2), (2,))

  374.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  375. 

  376.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  377.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  378.     b = Tensor(np.ones([64]), dtype=ms.float32)

  379.     compile_net(net, x, y, b)

  380. 

  381. 

  382. def test_matmul_add_broadcast2():

  383.     """

  384.     Feature: distribute operator sub in auto parallel.

  385.     Description: matmul-add broadcast net with strategy in semi auto parallel.

  386.     Expectation: compile done without error.

  387.     """

  388.     class Net(nn.Cell):

  389.         def __init__(self, strategy1, strategy2):

  390.             super().__init__()

  391.             self.matmul = P.MatMul().shard(strategy1)

  392.             self.add = P.Add().shard(strategy2)

  393. 

  394.         def construct(self, x, y, b):

  395.             out = self.matmul(x, y)

  396.             out = self.add(out, b)

  397.             return out

  398. 

  399.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  400.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  401.     strategy1 = ((2, 4), (4, 1))

  402.     strategy2 = ((4, 1), (1, 2))

  403.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  404. 

  405.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  406.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  407.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  408.     compile_net(net, x, y, b)

  409. 

  410. 

  411. def test_matmul_sub_broadcast():

  412.     """

  413.     Feature: distribute operator sub in auto parallel.

  414.     Description: matmul-sub broadcast net with strategy in semi auto parallel.

  415.     Expectation: compile done without error.

  416.     """

  417.     class Net(nn.Cell):

  418.         def __init__(self, strategy1, strategy2):

  419.             super().__init__()

  420.             self.matmul = P.MatMul().shard(strategy1)

  421.             self.sub = P.Sub().shard(strategy2)

  422. 

  423.         def construct(self, x, y, b):

  424.             out = self.matmul(x, y)

  425.             out = self.sub(out, b)

  426.             return out

  427. 

  428.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  429.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  430.     strategy1 = ((2, 2), (2, 2))

  431.     strategy2 = ((4, 2), (2,))

  432.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  433. 

  434.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  435.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  436.     b = Tensor(np.ones([64]), dtype=ms.float32)

  437.     compile_net(net, x, y, b)

  438. 

  439. 

  440. def test_matmul_sub_broadcast2():

  441.     """

  442.     Feature: distribute operator sub in auto parallel.

  443.     Description: matmul-sub broadcast net with strategy in semi auto parallel.

  444.     Expectation: compile done without error.

  445.     """

  446.     class Net(nn.Cell):

  447.         def __init__(self, strategy1, strategy2):

  448.             super().__init__()

  449.             self.matmul = P.MatMul().shard(strategy1)

  450.             self.sub = P.Sub().shard(strategy2)

  451. 

  452.         def construct(self, x, y, b):

  453.             out = self.matmul(x, y)

  454.             out = self.sub(out, b)

  455.             return out

  456. 

  457.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  458.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  459.     strategy1 = ((2, 4), (4, 1))

  460.     strategy2 = ((4, 1), (1, 2))

  461.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  462. 

  463.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  464.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  465.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  466.     compile_net(net, x, y, b)

  467. 

  468. 

  469. def test_matmul_mul_broadcast():

  470.     """

  471.     Feature: distribute operator sub in auto parallel.

  472.     Description: matmul-mul broadcast net with strategy in semi auto parallel.

  473.     Expectation: compile done without error.

  474.     """

  475.     class Net(nn.Cell):

  476.         def __init__(self, strategy1, strategy2):

  477.             super().__init__()

  478.             self.matmul = P.MatMul().shard(strategy1)

  479.             self.mul = P.Mul().shard(strategy2)

  480. 

  481.         def construct(self, x, y, b):

  482.             out = self.matmul(x, y)

  483.             out = self.mul(out, b)

  484.             return out

  485. 

  486.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  487.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  488.     strategy1 = ((2, 2), (2, 2))

  489.     strategy2 = ((4, 2), (2,))

  490.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  491. 

  492.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  493.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  494.     b = Tensor(np.ones([64]), dtype=ms.float32)

  495.     compile_net(net, x, y, b)

  496. 

  497. 

  498. def test_matmul_mul_broadcast2():

  499.     """

  500.     Feature: distribute operator sub in auto parallel.

  501.     Description: matmul-mul broadcast net with strategy in semi auto parallel.

  502.     Expectation: compile done without error.

  503.     """

  504.     class Net(nn.Cell):

  505.         def __init__(self, strategy1, strategy2):

  506.             super().__init__()

  507.             self.matmul = P.MatMul().shard(strategy1)

  508.             self.mul = P.Mul().shard(strategy2)

  509. 

  510.         def construct(self, x, y, b):

  511.             out = self.matmul(x, y)

  512.             out = self.mul(out, b)

  513.             return out

  514. 

  515.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  516.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  517.     strategy1 = ((2, 4), (4, 1))

  518.     strategy2 = ((4, 1), (1, 2))

  519.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  520. 

  521.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  522.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  523.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  524.     compile_net(net, x, y, b)

  525. 

  526. 

  527. def test_matmul_div_broadcast():

  528.     """

  529.     Feature: distribute operator sub in auto parallel.

  530.     Description: matmul-div broadcast net with strategy in semi auto parallel.

  531.     Expectation: compile done without error.

  532.     """

  533.     class Net(nn.Cell):

  534.         def __init__(self, strategy1, strategy2):

  535.             super().__init__()

  536.             self.matmul = P.MatMul().shard(strategy1)

  537.             self.div = P.Div().shard(strategy2)

  538. 

  539.         def construct(self, x, y, b):

  540.             out = self.matmul(x, y)

  541.             out = self.div(out, b)

  542.             return out

  543. 

  544.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  545.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  546.     strategy1 = ((2, 2), (2, 2))

  547.     strategy2 = ((4, 2), (2,))

  548.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  549. 

  550.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  551.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  552.     b = Tensor(np.ones([64]), dtype=ms.float32)

  553.     compile_net(net, x, y, b)

  554. 

  555. 

  556. def test_matmul_div_broadcast2():

  557.     """

  558.     Feature: distribute operator sub in auto parallel.

  559.     Description: matmul-div broadcast net with strategy in semi auto parallel.

  560.     Expectation: compile done without error.

  561.     """

  562.     class Net(nn.Cell):

  563.         def __init__(self, strategy1, strategy2):

  564.             super().__init__()

  565.             self.matmul = P.MatMul().shard(strategy1)

  566.             self.div = P.Div().shard(strategy2)

  567. 

  568.         def construct(self, x, y, b):

  569.             out = self.matmul(x, y)

  570.             out = self.div(out, b)

  571.             return out

  572. 

  573.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  574.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  575.     strategy1 = ((2, 4), (4, 1))

  576.     strategy2 = ((4, 1), (1, 2))

  577.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  578. 

  579.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  580.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  581.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  582.     compile_net(net, x, y, b)

  583. 

  584. 

  585. def test_matmul_greater_broadcast():

  586.     """

  587.     Feature: distribute operator sub in auto parallel.

  588.     Description: matmul-greater broadcast net with strategy in semi auto parallel.

  589.     Expectation: compile done without error.

  590.     """

  591.     class Net(nn.Cell):

  592.         def __init__(self, strategy1, strategy2):

  593.             super().__init__()

  594.             self.matmul = P.MatMul().shard(strategy1)

  595.             self.greater = P.Greater().shard(strategy2)

  596. 

  597.         def construct(self, x, y, b):

  598.             out = self.matmul(x, y)

  599.             out = self.greater(out, b)

  600.             return out

  601. 

  602.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  603.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  604.     strategy1 = ((2, 2), (2, 2))

  605.     strategy2 = ((4, 2), (2,))

  606.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  607. 

  608.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  609.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  610.     b = Tensor(np.ones([64]), dtype=ms.float32)

  611.     compile_net(net, x, y, b)

  612. 

  613. 

  614. def test_matmul_greater_broadcast2():

  615.     """

  616.     Feature: distribute operator sub in auto parallel.

  617.     Description: matmul-greater broadcast net with strategy in semi auto parallel.

  618.     Expectation: compile done without error.

  619.     """

  620.     class Net(nn.Cell):

  621.         def __init__(self, strategy1, strategy2):

  622.             super().__init__()

  623.             self.matmul = P.MatMul().shard(strategy1)

  624.             self.greater = P.Greater().shard(strategy2)

  625. 

  626.         def construct(self, x, y, b):

  627.             out = self.matmul(x, y)

  628.             out = self.greater(out, b)

  629.             return out

  630. 

  631.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  632.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  633.     strategy1 = ((2, 4), (4, 1))

  634.     strategy2 = ((4, 1), (1, 2))

  635.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  636. 

  637.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  638.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  639.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  640.     compile_net(net, x, y, b)

  641. 

  642. 

  643. def test_matmul_floordiv():

  644.     """

  645.     Feature: distribute operator sub in auto parallel.

  646.     Description: matmul-floordiv net with strategy in semi auto parallel.

  647.     Expectation: compile done without error.

  648.     """

  649.     class Net(nn.Cell):

  650.         def __init__(self, strategy1, strategy2):

  651.             super().__init__()

  652.             self.matmul = P.MatMul().shard(strategy1)

  653.             self.floordiv = P.FloorDiv().shard(strategy2)

  654. 

  655.         def construct(self, x, y, b):

  656.             out = self.matmul(x, y)

  657.             out = self.floordiv(out, b)

  658.             return out

  659. 

  660.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  661.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  662.     strategy1 = ((2, 2), (2, 2))

  663.     strategy2 = ((4, 2), (4, 2))

  664.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  665. 

  666.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  667.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  668.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  669.     compile_net(net, x, y, b)

  670. 

  671. 

  672. def test_matmul_floordiv_broadcast():

  673.     """

  674.     Feature: distribute operator sub in auto parallel.

  675.     Description: matmul-floordiv broadcast net with strategy in semi auto parallel.

  676.     Expectation: compile done without error.

  677.     """

  678.     class Net(nn.Cell):

  679.         def __init__(self, strategy1, strategy2):

  680.             super().__init__()

  681.             self.matmul = P.MatMul().shard(strategy1)

  682.             self.floordiv = P.FloorDiv().shard(strategy2)

  683. 

  684.         def construct(self, x, y, b):

  685.             out = self.matmul(x, y)

  686.             out = self.floordiv(out, b)

  687.             return out

  688. 

  689.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  690.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  691.     strategy1 = ((2, 2), (2, 2))

  692.     strategy2 = ((4, 2), (2,))

  693.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  694. 

  695.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  696.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  697.     b = Tensor(np.ones([64]), dtype=ms.float32)

  698.     compile_net(net, x, y, b)

  699. 

  700. 

  701. def test_matmul_floordiv_broadcast2():

  702.     """

  703.     Feature: distribute operator sub in auto parallel.

  704.     Description: matmul-floordiv broadcast net with strategy in semi auto parallel.

  705.     Expectation: compile done without error.

  706.     """

  707.     class Net(nn.Cell):

  708.         def __init__(self, strategy1, strategy2):

  709.             super().__init__()

  710.             self.matmul = P.MatMul().shard(strategy1)

  711.             self.floordiv = P.FloorDiv().shard(strategy2)

  712. 

  713.         def construct(self, x, y, b):

  714.             out = self.matmul(x, y)

  715.             out = self.floordiv(out, b)

  716.             return out

  717. 

  718.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  719.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  720.     strategy1 = ((2, 4), (4, 1))

  721.     strategy2 = ((4, 1), (1, 2))

  722.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  723. 

  724.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  725.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  726.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  727.     compile_net(net, x, y, b)

  728. 

  729. 

  730. def test_assign_sub():

  731.     """

  732.     Feature: distribute operator sub in auto parallel.

  733.     Description: mul-assign_sub net with strategy in semi auto parallel.

  734.     Expectation: compile done without error.

  735.     """

  736.     class Net(nn.Cell):

  737.         def __init__(self):

  738.             super().__init__()

  739.             self.assign_sub = P.AssignSub()

  740.             self.mul = P.Mul()

  741.             self.mul_weight = Parameter(Tensor(np.full([128, 32],

  742.                                                        0.5, dtype=np.float32)),

  743.                                         name="mul_weight")

  744.             self.assignsub_weight = Parameter(Tensor(np.full([128, 32],

  745.                                                              1.1, dtype=np.float32)),

  746.                                               name="assignsub_weight")

  747. 

  748.         def construct(self, x):

  749.             out = self.mul(x, self.mul_weight)

  750.             out = self.assign_sub(self.assignsub_weight, out)

  751.             return out

  752. 

  753.     class SubNetWithLoss(nn.Cell):

  754.         def __init__(self, network):

  755.             super(SubNetWithLoss, self).__init__()

  756.             self.loss = VirtualLoss()

  757.             self.network = network

  758. 

  759.         def construct(self, x):

  760.             predict = self.network(x,)

  761.             return self.loss(predict)

  762. 

  763.     class SubGradWrap(nn.Cell):

  764.         def __init__(self, network):

  765.             super(SubGradWrap, self).__init__()

  766.             self.network = network

  767. 

  768.         def construct(self, x):

  769.             return grad_all(self.network)(x)

  770. 

  771.     def compile_sub_net(net, x):

  772.         net.set_auto_parallel()

  773.         net.set_train()

  774.         _cell_graph_executor.compile(net, x)

  775. 

  776.     context.set_auto_parallel_context(device_num=64, global_rank=15)

  777.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  778.     net = SubGradWrap(SubNetWithLoss(Net()))

  779.     x = Tensor(np.ones([128, 32]), dtype=ms.float32)

  780.     compile_sub_net(net, x)

  781. 

  782. 

  783. def test_assign_add():

  784.     """

  785.     Feature: distribute operator sub in auto parallel.

  786.     Description: mul-assign_add net with strategy in semi auto parallel.

  787.     Expectation: compile done without error.

  788.     """

  789.     class Net(nn.Cell):

  790.         def __init__(self):

  791.             super().__init__()

  792.             self.assign_sub = P.AssignAdd()

  793.             self.mul = P.Mul()

  794.             self.mul_weight = Parameter(Tensor(np.full([128, 32],

  795.                                                        0.5, dtype=np.float32)),

  796.                                         name="mul_weight")

  797.             self.assignsub_weight = Parameter(Tensor(np.full([128, 32],

  798.                                                              1.1, dtype=np.float32)),

  799.                                               name="assignsub_weight")

  800. 

  801.         def construct(self, x):

  802.             out = self.mul(x, self.mul_weight)

  803.             out = self.assign_sub(self.assignsub_weight, out)

  804.             return out

  805. 

  806.     class SubNetWithLoss(nn.Cell):

  807.         def __init__(self, network):

  808.             super(SubNetWithLoss, self).__init__()

  809.             self.loss = VirtualLoss()

  810.             self.network = network

  811. 

  812.         def construct(self, x):

  813.             predict = self.network(x,)

  814.             return self.loss(predict)

  815. 

  816.     class SubGradWrap(nn.Cell):

  817.         def __init__(self, network):

  818.             super(SubGradWrap, self).__init__()

  819.             self.network = network

  820. 

  821.         def construct(self, x):

  822.             return grad_all(self.network)(x)

  823. 

  824.     def compile_sub_net(net, x):

  825.         net.set_auto_parallel()

  826.         net.set_train()

  827.         _cell_graph_executor.compile(net, x)

  828. 

  829.     context.set_auto_parallel_context(device_num=64, global_rank=15)

  830.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  831.     net = SubGradWrap(SubNetWithLoss(Net()))

  832.     x = Tensor(np.ones([128, 32]), dtype=ms.float32)

  833.     compile_sub_net(net, x)

  834. 

  835. 

  836. def test_assign():

  837.     """

  838.     Feature: distribute operator sub in auto parallel.

  839.     Description: mul-assign_sub net with strategy in semi auto parallel.

  840.     Expectation: compile done without error.

  841.     """

  842.     class Net(nn.Cell):

  843.         def __init__(self):

  844.             super().__init__()

  845.             self.assign_sub = P.Assign()

  846.             self.mul = P.Mul()

  847.             self.mul_weight = Parameter(Tensor(np.full([128, 32],

  848.                                                        0.5, dtype=np.float32)),

  849.                                         name="mul_weight")

  850.             self.assignsub_weight = Parameter(Tensor(np.full([128, 32],

  851.                                                              1.1, dtype=np.float32)),

  852.                                               name="assignsub_weight")

  853. 

  854.         def construct(self, x):

  855.             out = self.mul(x, self.mul_weight)

  856.             out = self.assign_sub(self.assignsub_weight, out)

  857.             return out

  858. 

  859.     class SubNetWithLoss(nn.Cell):

  860.         def __init__(self, network):

  861.             super(SubNetWithLoss, self).__init__()

  862.             self.loss = VirtualLoss()

  863.             self.network = network

  864. 

  865.         def construct(self, x):

  866.             predict = self.network(x,)

  867.             return self.loss(predict)

  868. 

  869.     class SubGradWrap(nn.Cell):

  870.         def __init__(self, network):

  871.             super(SubGradWrap, self).__init__()

  872.             self.network = network

  873. 

  874.         def construct(self, x):

  875.             return grad_all(self.network)(x)

  876. 

  877.     def compile_sub_net(net, x):

  878.         net.set_auto_parallel()

  879.         net.set_train()

  880.         _cell_graph_executor.compile(net, x)

  881. 

  882.     context.set_auto_parallel_context(device_num=64, global_rank=15)

  883.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  884.     net = SubGradWrap(SubNetWithLoss(Net()))

  885.     x = Tensor(np.ones([128, 32]), dtype=ms.float32)

  886.     compile_sub_net(net, x)

  887. 

  888. 

  889. def test_matmul_bitwise_and_broadcast():

  890.     """

  891.     Feature: distribute operator BitwiseAnd in auto parallel.

  892.     Description: mul-BitwiseAnd net with strategy in semi auto parallel.

  893.     Expectation: compile done without error.

  894.     """

  895.     class Net(nn.Cell):

  896.         def __init__(self, strategy1, strategy2):

  897.             super().__init__()

  898.             self.bitwise_and = P.BitwiseAnd().shard(strategy1)

  899.             self.matmul = P.MatMul().shard(strategy2)

  900. 

  901. 

  902.         def construct(self, x, y, z):

  903.             out = self.bitwise_and(x, y)

  904.             out = self.matmul(out, z)

  905. 

  906.             return out

  907. 

  908.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  909.     strategy1 = ((2, 1), (1, 4))

  910.     strategy2 = ((1, 4), (4, 2))

  911.     net = Net(strategy1, strategy2)

  912. 

  913.     x = Tensor(np.ones([64, 1]), dtype=ms.int32)

  914.     y = Tensor(np.ones([1, 64]), dtype=ms.int32)

  915.     z = Tensor(np.ones([64, 32]), dtype=ms.int32)

  916.     compile_net(net, x, y, z)

  917. 

  918. 

  919. def test_matmul_bitwise_or_broadcast():

  920.     """

  921.     Feature: distribute operator BitwiseOr in auto parallel.

  922.     Description: mul-BitwiseOr net with strategy in semi auto parallel.

  923.     Expectation: compile done without error.

  924.     """

  925.     class Net(nn.Cell):

  926.         def __init__(self, strategy1, strategy2):

  927.             super().__init__()

  928.             self.bitwise_or = P.BitwiseOr().shard(strategy1)

  929.             self.matmul = P.MatMul().shard(strategy2)

  930. 

  931.         def construct(self, x, y, z):

  932.             out = self.bitwise_or(x, y)

  933.             out = self.matmul(out, z)

  934.             return out

  935. 

  936.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  937.     strategy1 = ((2, 1), (1, 4))

  938.     strategy2 = ((1, 4), (4, 2))

  939.     net = Net(strategy1, strategy2)

  940. 

  941.     x = Tensor(np.ones([64, 1]), dtype=ms.int32)

  942.     y = Tensor(np.ones([1, 64]), dtype=ms.int32)

  943.     z = Tensor(np.ones([64, 32]), dtype=ms.int32)

  944.     compile_net(net, x, y, z)

  945. 

  946. 

  947. def test_matmul_bitwise_xor_broadcast():

  948.     """

  949.     Feature: distribute operator BitwiseXor in auto parallel.

  950.     Description: mul-BitwiseXor net with strategy in semi auto parallel.

  951.     Expectation: compile done without error.

  952.     """

  953.     class Net(nn.Cell):

  954.         def __init__(self, strategy1, strategy2):

  955.             super().__init__()

  956.             self.bitwise_xor = P.BitwiseXor().shard(strategy1)

  957.             self.matmul = P.MatMul().shard(strategy2)

  958. 

  959.         def construct(self, x, y, z):

  960.             out = self.bitwise_xor(x, y)

  961.             out = self.matmul(out, z)

  962.             return out

  963. 

  964.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  965.     strategy1 = ((2, 1), (1, 4))

  966.     strategy2 = ((1, 4), (4, 2))

  967.     net = Net(strategy1, strategy2)

  968. 

  969.     x = Tensor(np.ones([64, 1]), dtype=ms.int32)

  970.     y = Tensor(np.ones([1, 64]), dtype=ms.int32)

  971.     z = Tensor(np.ones([64, 32]), dtype=ms.int32)

  972.     compile_net(net, x, y, z)

  973. 

  974. 

  975. def test_matmul_mul_no_nan_broadcast():

  976.     """

  977.     Feature: distribute operator MulNoNan in auto parallel.

  978.     Description: mul-MulNoNan net with strategy in semi auto parallel.

  979.     Expectation: compile done without error.

  980.     """

  981.     class Net(nn.Cell):

  982.         def __init__(self, strategy1, strategy2):

  983.             super().__init__()

  984.             self.matmul = P.MatMul().shard(strategy1)

  985.             self.mul_no_nan = P.MulNoNan().shard(strategy2)

  986. 

  987.         def construct(self, x, y, b):

  988.             out = self.matmul(x, y)

  989.             out = self.mul_no_nan(out, b)

  990.             return out

  991. 

  992.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  993.     strategy1 = ((2, 4), (4, 1))

  994.     strategy2 = ((4, 1), (1, 2))

  995.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  996. 

  997.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  998.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  999.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  1000.     compile_net(net, x, y, b)

  1001. 

  1002. 

  1003. def test_matmul_truncate_div_broadcast():

  1004.     """

  1005.     Feature: distribute operator TruncateDiv in auto parallel.

  1006.     Description: mul-TruncateDiv net with strategy in semi auto parallel.

  1007.     Expectation: compile done without error.

  1008.     """

  1009.     class Net(nn.Cell):

  1010.         def __init__(self, strategy1, strategy2):

  1011.             super().__init__()

  1012.             self.matmul = P.MatMul().shard(strategy1)

  1013.             self.truncate_div = P.TruncateDiv().shard(strategy2)

  1014. 

  1015.         def construct(self, x, y, b):

  1016.             out = self.matmul(x, y)

  1017.             out = self.truncate_div(out, b)

  1018.             return out

  1019. 

  1020.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  1021.     strategy1 = ((2, 4), (4, 1))

  1022.     strategy2 = ((4, 1), (1, 2))

  1023.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  1024. 

  1025.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  1026.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  1027.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  1028.     compile_net(net, x, y, b)

  1029. 

  1030. 

  1031. def test_matmul_truncate_mod_broadcast():

  1032.     """

  1033.     Feature: distribute operator TruncateMod in auto parallel.

  1034.     Description: mul-TruncateMod net with strategy in semi auto parallel.

  1035.     Expectation: compile done without error.

  1036.     """

  1037.     class Net(nn.Cell):

  1038.         def __init__(self, strategy1, strategy2):

  1039.             super().__init__()

  1040.             self.matmul = P.MatMul().shard(strategy1)

  1041.             self.truncate_mod = P.TruncateMod().shard(strategy2)

  1042. 

  1043.         def construct(self, x, y, b):

  1044.             out = self.matmul(x, y)

  1045.             out = self.truncate_mod(out, b)

  1046.             return out

  1047. 

  1048.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  1049.     strategy1 = ((2, 4), (4, 1))

  1050.     strategy2 = ((4, 1), (1, 2))

  1051.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  1052. 

  1053.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  1054.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  1055.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  1056.     compile_net(net, x, y, b)

  1057. 

  1058. 

  1059. def test_matmul_xdivy_broadcast():

  1060.     """

  1061.     Feature: distribute operator Xdivy in auto parallel.

  1062.     Description: mul-Xdivy net with strategy in semi auto parallel.

  1063.     Expectation: compile done without error.

  1064.     """

  1065.     class Net(nn.Cell):

  1066.         def __init__(self, strategy1, strategy2):

  1067.             super().__init__()

  1068.             self.matmul = P.MatMul().shard(strategy1)

  1069.             self.xdivy = P.Xdivy().shard(strategy2)

  1070. 

  1071.         def construct(self, x, y, b):

  1072.             out = self.matmul(x, y)

  1073.             out = self.xdivy(out, b)

  1074.             return out

  1075. 

  1076.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  1077.     strategy1 = ((2, 4), (4, 1))

  1078.     strategy2 = ((4, 1), (1, 2))

  1079.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  1080. 

  1081.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  1082.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  1083.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  1084.     compile_net(net, x, y, b)

  1085. 

  1086. 

  1087. def test_matmul_xlogy_broadcast():

  1088.     """

  1089.     Feature: distribute operator Xlogy in auto parallel.

  1090.     Description: mul-Xlogy net with strategy in semi auto parallel.

  1091.     Expectation: compile done without error.

  1092.     """

  1093.     class Net(nn.Cell):

  1094.         def __init__(self, strategy1, strategy2):

  1095.             super().__init__()

  1096.             self.matmul = P.MatMul().shard(strategy1)

  1097.             self.xlogy = P.Xlogy().shard(strategy2)

  1098. 

  1099.         def construct(self, x, y, b):

  1100.             out = self.matmul(x, y)

  1101.             out = self.xlogy(out, b)

  1102.             return out

  1103. 

  1104.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  1105.     strategy1 = ((2, 4), (4, 1))

  1106.     strategy2 = ((4, 1), (1, 2))

  1107.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  1108. 

  1109.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  1110.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  1111.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  1112.     compile_net(net, x, y, b)

  1113. 

  1114. 

  1115. def test_matmul_squared_difference_broadcast():

  1116.     """

  1117.     Feature: distribute operator SquaredDifference in auto parallel.

  1118.     Description: mul-SquaredDifference net with strategy in semi auto parallel.

  1119.     Expectation: compile done without error.

  1120.     """

  1121.     class Net(nn.Cell):

  1122.         def __init__(self, strategy1, strategy2):

  1123.             super().__init__()

  1124.             self.matmul = P.MatMul().shard(strategy1)

  1125.             self.squared_difference = P.SquaredDifference().shard(strategy2)

  1126. 

  1127.         def construct(self, x, y, b):

  1128.             out = self.matmul(x, y)

  1129.             out = self.squared_difference(out, b)

  1130.             return out

  1131. 

  1132.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  1133.     strategy1 = ((2, 4), (4, 1))

  1134.     strategy2 = ((4, 1), (1, 2))

  1135.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  1136. 

  1137.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  1138.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  1139.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  1140.     compile_net(net, x, y, b)

  1141. 

  1142. 

  1143. def test_matmul_masked_fill_broadcast_with_value_float():

  1144.     """

  1145.     Feature: distribute operator MaskedFill in auto parallel.

  1146.     Description: mul-MaskedFill net with strategy in semi auto parallel.

  1147.     Expectation: compile done without error.

  1148.     """

  1149.     class Net(nn.Cell):

  1150.         def __init__(self, strategy1, strategy2):

  1151.             super().__init__()

  1152.             self.matmul = P.MatMul().shard(strategy1)

  1153.             self.masked_fill = P.MaskedFill().shard(strategy2)

  1154.             self.value = 1.0

  1155. 

  1156.         def construct(self, x, y, b):

  1157.             out = self.matmul(x, y)

  1158.             out = self.masked_fill(out, b, self.value)

  1159.             return out

  1160. 

  1161.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  1162.     strategy1 = ((2, 4), (4, 1))

  1163.     strategy2 = ((4, 1), (1, 2))

  1164.     net = Net(strategy1, strategy2)

  1165. 

  1166.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  1167.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  1168.     b = Tensor(np.ones([1, 64]), dtype=ms.bool_)

  1169.     compile_net(net, x, y, b)

  1170. 

  1171. 

  1172. def test_matmul_masked_fill_broadcast_with_value_tensor():

  1173.     """

  1174.     Feature: distribute operator MaskedFill in auto parallel.

  1175.     Description: mul-MaskedFill net with strategy in semi auto parallel.

  1176.     Expectation: compile done without error.

  1177.     """

  1178.     class Net(nn.Cell):

  1179.         def __init__(self, strategy1, strategy2):

  1180.             super().__init__()

  1181.             self.matmul = P.MatMul().shard(strategy1)

  1182.             self.masked_fill = P.MaskedFill().shard(strategy2)

  1183.             self.value = Tensor(1.0, ms.float32)

  1184. 

  1185.         def construct(self, x, y, b):

  1186.             out = self.matmul(x, y)

  1187.             out = self.masked_fill(out, b, self.value)

  1188.             return out

  1189. 

  1190.     context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="semi_auto_parallel")

  1191.     strategy1 = ((2, 4), (4, 1))

  1192.     strategy2 = ((4, 1), (1, 2), ())

  1193.     net = Net(strategy1, strategy2)

  1194. 

  1195.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  1196.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  1197.     b = Tensor(np.ones([1, 64]), dtype=ms.bool_)

  1198.     compile_net(net, x, y, b)