• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import pytest
2import numpy as np
3
4import mindspore.context as context
5import mindspore.nn as nn
6from mindspore import Tensor
7from mindspore.common.api import ms_function
8from mindspore.ops import operations as P
9
10
11class Net(nn.Cell):
12    def __init__(self, seq_dim, batch_dim):
13        super(Net, self).__init__()
14        self.reverse_sequence = P.ReverseSequence(
15            seq_dim=seq_dim, batch_dim=batch_dim)
16
17    @ms_function
18    def construct(self, x, seq_lengths):
19        return self.reverse_sequence(x, seq_lengths)
20
21
22@pytest.mark.level0
23@pytest.mark.platform_x86_gpu_training
24@pytest.mark.env_onecard
25def test_net_int8():
26    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
27    x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).astype(np.int8)
28    seq_lengths = np.array([1, 2, 3]).astype(np.int32)
29    seq_dim = 0
30    batch_dim = 1
31    net = Net(seq_dim, batch_dim)
32    output = net(Tensor(x), Tensor(seq_lengths))
33    expected = np.array([[1, 5, 9], [4, 2, 6], [7, 8, 3]]).astype(np.int8)
34    assert np.array_equal(output.asnumpy(), expected)
35
36
37@pytest.mark.level0
38@pytest.mark.platform_x86_gpu_training
39@pytest.mark.env_onecard
40def test_net_int32():
41    context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
42    x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).astype(np.int32)
43    seq_lengths = np.array([1, 2, 3]).astype(np.int64)
44    seq_dim = 1
45    batch_dim = 0
46    net = Net(seq_dim, batch_dim)
47    output = net(Tensor(x), Tensor(seq_lengths))
48    expected = np.array([[1, 2, 3], [5, 4, 6], [9, 8, 7]]).astype(np.int32)
49    assert np.array_equal(output.asnumpy(), expected)
50
51
52@pytest.mark.level0
53@pytest.mark.platform_x86_gpu_training
54@pytest.mark.env_onecard
55def test_net_float32():
56    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
57    x = np.array([[[[1, 2], [3, 4]],
58                   [[5, 6], [7, 8]],
59                   [[9, 10], [11, 12]],
60                   [[13, 14], [15, 16]]],
61                  [[[17, 18], [19, 20]],
62                   [[21, 22], [23, 24]],
63                   [[25, 26], [27, 28]],
64                   [[29, 30], [31, 21]]]]).astype(np.float32)
65    seq_lengths = np.array([2, 2, 2, 2]).astype(np.int64)
66    seq_dim = 0
67    batch_dim = 1
68    net = Net(seq_dim, batch_dim)
69    output = net(Tensor(x), Tensor(seq_lengths))
70    expected = np.array([[[[17., 18.], [19., 20.]],
71                          [[21., 22.], [23., 24.]],
72                          [[25., 26.], [27., 28.]],
73                          [[29., 30.], [31., 21.]]],
74                         [[[1., 2.], [3., 4.]],
75                          [[5., 6.], [7., 8.]],
76                          [[9., 10.], [11., 12.]],
77                          [[13., 14.], [15., 16.]]]]).astype(np.float32)
78    assert np.array_equal(output.asnumpy(), expected)
79
80
81@pytest.mark.level0
82@pytest.mark.platform_x86_gpu_training
83@pytest.mark.env_onecard
84def test_net_float64_0_dim():
85    """
86    Test added to test for 0 seq len edge case
87    """
88    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
89    x = np.array([[[[1, 2], [3, 4]],
90                   [[5, 6], [7, 8]],
91                   [[9, 10], [11, 12]],
92                   [[13, 14], [15, 16]]],
93                  [[[17, 18], [19, 20]],
94                   [[21, 22], [23, 24]],
95                   [[25, 26], [27, 28]],
96                   [[29, 30], [31, 21]]]]).astype(np.float32)
97    seq_lengths = np.array([2, 2, 0, 0]).astype(np.int64)
98    seq_dim = 2
99    batch_dim = 1
100    net = Net(seq_dim, batch_dim)
101    output = net(Tensor(x), Tensor(seq_lengths))
102    expected = np.array([[[[3., 4.], [1., 2.]],
103                          [[7.,  8.], [5., 6.]],
104                          [[9., 10.], [11., 12.]],
105                          [[13., 14.], [15., 16.]]],
106                         [[[19., 20.], [17., 18.]],
107                          [[23., 24.], [21., 22.]],
108                          [[25., 26.], [27., 28.]],
109                          [[29., 30.], [31., 21.]]]]).astype(np.float32)
110    assert np.array_equal(output.asnumpy(), expected)
111