# Copyright 2020 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ import numpy as np import pytest import mindspore.context as context import mindspore.nn as nn from mindspore import Tensor from mindspore.common import dtype as mstype from mindspore.ops import operations as P class BatchMatMulNet(nn.Cell): def __init__(self, transpose_a=False, transpose_b=False): super(BatchMatMulNet, self).__init__() self.batch_matmul = P.BatchMatMul(transpose_a, transpose_b) def construct(self, x, y): return self.batch_matmul(x, y) def judge_result_correct(result, expect): assert result.dtype == expect.dtype assert result.shape == expect.shape assert np.allclose(result, expect) @pytest.mark.level0 @pytest.mark.platform_x86_cpu @pytest.mark.env_onecard def test_4d_no_transpose_vec(): input_x = Tensor(np.arange(2 * 4 * 1 * 3).reshape((2, 4, 1, 3)), mstype.float32) input_y = Tensor(np.arange(2 * 4 * 3 * 4).reshape((2, 4, 3, 4)), mstype.float32) context.set_context(mode=context.GRAPH_MODE, device_target='CPU') net = BatchMatMulNet() output = net(input_x, input_y) expect = np.array([[[[20, 23, 26, 29]], [[200, 212, 224, 236]], [[596, 617, 638, 659]], [[1208, 1238, 1268, 1298]]], [[[2036, 2075, 2114, 2153]], [[3080, 3128, 3176, 3224]], [[4340, 4397, 4454, 4511]], [[5816, 5882, 5948, 6014]]]], dtype=np.float32) judge_result_correct(output.asnumpy(), expect) @pytest.mark.level0 @pytest.mark.platform_x86_cpu @pytest.mark.env_onecard def test_4d_no_transpose(): input_x = Tensor(np.arange(2 * 3 * 2 * 3).reshape((2, 3, 2, 3)), mstype.float32) input_y = Tensor(np.arange(2 * 3 * 3 * 4).reshape((2, 3, 3, 4)), mstype.float32) context.set_context(mode=context.GRAPH_MODE, device_target='CPU') net = BatchMatMulNet() output = net(input_x, input_y) expect = np.array([[[[20., 23., 26., 29.], [56., 68., 80., 92.]], [[344., 365., 386., 407.], [488., 518., 548., 578.]], [[1100., 1139., 1178., 1217.], [1352., 1400., 1448., 1496.]]], [[[2288., 2345., 2402., 2459.], [2648., 2714., 2780., 2846.]], [[3908., 3983., 4058., 4133.], [4376., 4460., 4544., 4628.]], [[5960., 6053., 6146., 6239.], [6536., 6638., 6740., 6842.]]]], dtype=np.float32) judge_result_correct(output.asnumpy(), expect) @pytest.mark.level0 @pytest.mark.platform_x86_cpu @pytest.mark.env_onecard def test_4d_transpose_a(): input_x = Tensor(np.arange(2 * 3 * 3 * 2).reshape((2, 3, 3, 2)), mstype.float32) input_y = Tensor(np.arange(2 * 3 * 3 * 4).reshape((2, 3, 3, 4)), mstype.float32) context.set_context(mode=context.GRAPH_MODE, device_target="CPU") net = BatchMatMulNet(transpose_a=True) output = net(input_x, input_y) expect = np.array([[[[40., 46., 52., 58.], [52., 61., 70., 79.]], [[400., 424., 448., 472.], [448., 475., 502., 529.]], [[1192., 1234., 1276., 1318.], [1276., 1321., 1366., 1411.]]], [[[2416., 2476., 2536., 2596.], [2536., 2599., 2662., 2725.]], [[4072., 4150., 4228., 4306.], [4228., 4309., 4390., 4471.]], [[6160., 6256., 6352., 6448.], [6352., 6451., 6550., 6649.]]]], dtype=np.float32) judge_result_correct(output.asnumpy(), expect) @pytest.mark.level0 @pytest.mark.platform_x86_cpu @pytest.mark.env_onecard def test_4d_transpose_b(): input_x = Tensor(np.arange(2 * 3 * 2 * 3).reshape((2, 3, 2, 3)), mstype.float32) input_y = Tensor(np.arange(2 * 3 * 4 * 3).reshape((2, 3, 4, 3)), mstype.float32) context.set_context(mode=context.GRAPH_MODE, device_target="CPU") net = BatchMatMulNet(transpose_b=True) output = net(input_x, input_y) expect = np.array([[[[5.000e+00, 1.400e+01, 2.300e+01, 3.200e+01], [1.400e+01, 5.000e+01, 8.600e+01, 1.220e+02]], [[2.750e+02, 3.380e+02, 4.010e+02, 4.640e+02], [3.920e+02, 4.820e+02, 5.720e+02, 6.620e+02]], [[9.770e+02, 1.094e+03, 1.211e+03, 1.328e+03], [1.202e+03, 1.346e+03, 1.490e+03, 1.634e+03]]], [[[2.111e+03, 2.282e+03, 2.453e+03, 2.624e+03], [2.444e+03, 2.642e+03, 2.840e+03, 3.038e+03]], [[3.677e+03, 3.902e+03, 4.127e+03, 4.352e+03], [4.118e+03, 4.370e+03, 4.622e+03, 4.874e+03]], [[5.675e+03, 5.954e+03, 6.233e+03, 6.512e+03], [6.224e+03, 6.530e+03, 6.836e+03, 7.142e+03]]]], dtype=np.float32) judge_result_correct(output.asnumpy(), expect) @pytest.mark.level0 @pytest.mark.platform_x86_cpu @pytest.mark.env_onecard def test_4d_transpose_ab(): input_x = Tensor(np.arange(2 * 3 * 3 * 2).reshape((2, 3, 3, 2)), mstype.float16) input_y = Tensor(np.arange(2 * 3 * 4 * 3).reshape((2, 3, 4, 3)), mstype.float16) context.set_context(mode=context.GRAPH_MODE, device_target="CPU") net = BatchMatMulNet(transpose_a=True, transpose_b=True) output = net(input_x, input_y) expect = np.array([[[[10., 28., 46., 64.], [13., 40., 67., 94.]], [[316., 388., 460., 532.], [355., 436., 517., 598.]], [[1054., 1180., 1306., 1432.], [1129., 1264., 1399., 1534.]]], [[[2224., 2404., 2584., 2764.], [2335., 2524., 2713., 2902.]], [[3826., 4060., 4294., 4528.], [3973., 4216., 4459., 4702.]], [[5860., 6148., 6436., 6724.], [6043., 6340., 6637., 6934.]]]], np.float16) judge_result_correct(output.asnumpy(), expect)