mindspore/tests/st/ops/cpu/test_batchdot_op.py

# 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 pytest
import numpy as np

import mindspore
from mindspore import Tensor
import mindspore.nn as nn
import mindspore.context as context
from mindspore.ops import composite as C


class NetBatchDot(nn.Cell):
    def __init__(self, axes):
        super(NetBatchDot, self).__init__()
        self.axes = axes

    def construct(self, x, y):
        return C.batch_dot(x, y, self.axes)


# Implementation with numpy in tensorflow
def _reference_batch_dot(x, y, axes):
    if isinstance(axes, int):
        axes = [axes, axes]
    elif isinstance(axes, tuple):
        axes = list(axes)
    if axes is None:
        if y.ndim == 2:
            axes = [x.ndim - 1, y.ndim - 1]
        else:
            axes = [x.ndim - 1, y.ndim - 2]
    if axes[0] < 0:
        axes[0] += x.ndim
    if axes[1] < 0:
        axes[1] += y.ndim
    result = []
    axes = [axes[0] - 1, axes[1] - 1]
    for xi, yi in zip(x, y):
        result.append(np.tensordot(xi, yi, axes))
    result = np.array(result)
    if result.ndim == 1:
        result = np.expand_dims(result, -1)
    return result


@pytest.mark.level0
@pytest.mark.platform_x86_cpu
@pytest.mark.env_onecard
def test_batch_dot_fp32():
    context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
    np.random.seed(12876)

    # case 1
    shape_x1 = (3, 12, 5, 2, 3)
    shape_x2 = (3, 1, 7, 3, 2)
    axes = (-1, -2)
    x1 = np.ones(shape=shape_x1).astype(np.float32)
    x2 = np.ones(shape=shape_x2).astype(np.float32)
    x1_tensor = Tensor(x1, dtype=mindspore.float32)
    x2_tensor = Tensor(x2, dtype=mindspore.float32)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)

    assert np.allclose(ms_result_np, tf_result)

    # case 2
    shape_x1 = (4, 3, 7, 5)
    shape_x2 = (4, 1, 7, 1)
    axes = 2
    x1 = np.random.random(shape_x1).astype(np.float32)
    x2 = np.random.random(shape_x2).astype(np.float32)
    x1_tensor = Tensor(x1, dtype=mindspore.float32)
    x2_tensor = Tensor(x2, dtype=mindspore.float32)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)

    assert np.allclose(ms_result_np, tf_result)

    # case 3
    shape_x1 = (18, 3, 5, 7)
    shape_x2 = (18, 1, 3, 7)
    axes = -1
    x1 = np.random.random(shape_x1).astype(np.float32)
    x2 = np.random.random(shape_x2).astype(np.float32)
    x1_tensor = Tensor(x1, dtype=mindspore.float32)
    x2_tensor = Tensor(x2, dtype=mindspore.float32)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)

    assert np.allclose(ms_result_np, tf_result)

    # case 4
    shape_x1 = (2, 11, 3, 9)
    shape_x2 = (2, 7, 9, 3)
    axes = None
    x1 = np.random.random(shape_x1).astype(np.float32)
    x2 = np.random.random(shape_x2).astype(np.float32)
    x1_tensor = Tensor(x1, dtype=mindspore.float32)
    x2_tensor = Tensor(x2, dtype=mindspore.float32)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)

    assert np.allclose(ms_result_np, tf_result)

    # case 5
    shape_x1 = (7, 5)
    shape_x2 = (7, 5)
    axes = None
    x1 = np.random.random(shape_x1).astype(np.float32)
    x2 = np.random.random(shape_x2).astype(np.float32)
    x1_tensor = Tensor(x1, dtype=mindspore.float32)
    x2_tensor = Tensor(x2, dtype=mindspore.float32)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)

    assert np.allclose(ms_result_np, tf_result)

    # case 6
    shape_x1 = (7, 3, 5)
    shape_x2 = (7, 5)
    axes = None
    x1 = np.random.random(shape_x1).astype(np.float32)
    x2 = np.random.random(shape_x2).astype(np.float32)
    x1_tensor = Tensor(x1, dtype=mindspore.float32)
    x2_tensor = Tensor(x2, dtype=mindspore.float32)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)

    assert np.allclose(ms_result_np, tf_result)

    # case 7
    shape_x1 = (7, 5)
    shape_x2 = (7, 5, 3)
    axes = None
    x1 = np.random.random(shape_x1).astype(np.float32)
    x2 = np.random.random(shape_x2).astype(np.float32)
    x1_tensor = Tensor(x1, dtype=mindspore.float32)
    x2_tensor = Tensor(x2, dtype=mindspore.float32)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)

    assert np.allclose(ms_result_np, tf_result)

    # case 8
    shape_x1 = (39, 6)
    shape_x2 = (39, 6)
    axes = -1
    x1 = np.random.random(shape_x1).astype(np.float32)
    x2 = np.random.random(shape_x2).astype(np.float32)
    x1_tensor = Tensor(x1, dtype=mindspore.float32)
    x2_tensor = Tensor(x2, dtype=mindspore.float32)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)

    assert np.allclose(ms_result_np, tf_result)

    # case 9
    shape_x1 = (21, 2, 3)
    shape_x2 = (21, 3, 2)
    axes = (-1, -2)
    x1 = np.ones(shape=shape_x1).astype(np.float32)
    x2 = np.ones(shape=shape_x2).astype(np.float32)
    x1_tensor = Tensor(x1, dtype=mindspore.float32)
    x2_tensor = Tensor(x2, dtype=mindspore.float32)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)
    assert np.allclose(ms_result_np, tf_result)

    # case 10
    shape_x1 = (4, 3, 2, 1, 7, 5)
    shape_x2 = (4, 5, 7, 1)
    axes = -2
    x1 = np.ones(shape=shape_x1).astype(np.float32)
    x2 = np.ones(shape=shape_x2).astype(np.float32)
    x1_tensor = Tensor(x1, dtype=mindspore.float32)
    x2_tensor = Tensor(x2, dtype=mindspore.float32)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)
    assert np.allclose(ms_result_np, tf_result)

    # case 10
    shape_x1 = (4, 3, 2, 1, 7, 5)
    shape_x2 = (4, 5, 7, 1)
    axes = -2
    x1 = np.ones(shape=shape_x1).astype(np.float16)
    x2 = np.ones(shape=shape_x2).astype(np.float16)
    x1_tensor = Tensor(x1, dtype=mindspore.float16)
    x2_tensor = Tensor(x2, dtype=mindspore.float16)

    network = NetBatchDot(axes)
    ms_result_np = network(x1_tensor, x2_tensor).asnumpy()
    tf_result = _reference_batch_dot(x1, x2, axes)
    assert np.allclose(ms_result_np, tf_result)