• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2024 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
16import pytest
17import numpy as np
18import mindspore as ms
19from mindspore import ops, jit, JitConfig
20from mindspore.mint import cumsum
21from tests.st.ops.dynamic_shape.test_op_utils import TEST_OP
22from tests.st.utils import test_utils
23
24def generate_random_input(shape, dim):
25    x = np.random.randn(*shape)
26    return x, np.cumsum(x, dim)
27
28
29def cumsum_func(x, dim):
30    return cumsum(x, dim)
31
32
33@test_utils.run_with_cell
34def cumsum_forward_func(x, dim):
35    return cumsum_func(x, dim)
36
37
38def cumsum_bwd_func(x, dim):
39    return ops.grad(cumsum_func, (0,))(x, dim)
40
41
42@test_utils.run_with_cell
43def cumsum_backward_func(x, dim):
44    return cumsum_bwd_func(x, dim)
45
46
47@pytest.mark.level0
48@pytest.mark.env_onecard
49@pytest.mark.platform_arm_ascend_training
50@pytest.mark.platform_x86_ascend_training
51@pytest.mark.parametrize('mode', ['pynative', 'KBK'])
52def test_cumsum_forward(mode):
53    """
54    Feature: Ops.
55    Description: test op cumsum.
56    Expectation: expect correct result.
57    """
58    test_shape = (2, 3, 4, 5)
59    dim1 = 2
60    x, expect = generate_random_input(test_shape, dim1)
61    if mode == 'pynative':
62        ms.set_context(mode=ms.PYNATIVE_MODE)
63        output = cumsum_forward_func(ms.Tensor(x), dim1)
64    else:
65        output = (jit(cumsum_forward_func, jit_config=JitConfig(jit_level="O0")))(ms.Tensor(x), dim1)
66    np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-4)
67
68
69@pytest.mark.level1
70@pytest.mark.platform_arm_ascend910b_training
71@pytest.mark.env_onecard
72@pytest.mark.parametrize('mode', ['pynative', 'KBK'])
73def test_cumsum_bfloat16(mode):
74    """
75    Feature: test ne functional API.
76    Description: testcase for ne functional API.
77    Expectation: the result match with expected result.
78    """
79    test_shape = (2, 3, 4)
80    dim1 = 1
81    x, expect = generate_random_input(test_shape, dim1)
82    if mode == 'pynative':
83        ms.set_context(mode=ms.PYNATIVE_MODE)
84        output = cumsum_forward_func(ms.Tensor(x), dim1)
85    else:
86        output = (jit(cumsum_forward_func, jit_config=JitConfig(jit_level="O0")))(ms.Tensor(x), dim1)
87    np.testing.assert_allclose(output.float().asnumpy(), expect, rtol=5e-3, atol=5e-3)
88
89
90@pytest.mark.level0
91@pytest.mark.env_onecard
92@pytest.mark.platform_arm_ascend_training
93@pytest.mark.platform_x86_ascend_training
94@pytest.mark.parametrize('mode', ['pynative', 'KBK'])
95def test_cumsum_backward(mode):
96    """
97    Feature: Ops.
98    Description: test op cumsum.
99    Expectation: expect correct result.
100    """
101    test_shape = (2, 3, 4, 5)
102    dim1 = 0
103    x, _ = generate_random_input(test_shape, dim1)
104    expect = np.flip(np.cumsum(np.flip(np.ones(test_shape), dim1), dim1), dim1)
105    if mode == 'pynative':
106        ms.set_context(mode=ms.PYNATIVE_MODE)
107        output = cumsum_backward_func(ms.Tensor(x), dim1)
108    else:
109        output = (jit(cumsum_backward_func, jit_config=JitConfig(jit_level="O0")))(ms.Tensor(x), dim1)
110    np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-4)
111
112
113
114@pytest.mark.level0
115@pytest.mark.env_onecard
116@pytest.mark.platform_arm_ascend_training
117@pytest.mark.platform_x86_ascend_training
118def test_cumsum_dynamic_shape():
119    """
120    Feature: Test dynamic shape.
121    Description: test function div dynamic feature.
122    Expectation: expect correct result.
123    """
124    dim1 = 0
125    ms_data1, _ = generate_random_input((2, 3, 4), dim1)
126    dim2 = 1
127    ms_data2, _ = generate_random_input((3, 4, 5, 6), dim2)
128    TEST_OP(cumsum_forward_func, [[ms.Tensor(ms_data1), dim1], [ms.Tensor(ms_data2), dim2]],
129            '', disable_yaml_check=True, disable_mode=['GRAPH_MODE'])
130