• 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# ============================================================================
15import pytest
16import numpy as np
17import mindspore as ms
18from mindspore import ops, jit, JitConfig
19from tests.st.utils import test_utils
20from tests.st.ops.dynamic_shape.test_op_utils import TEST_OP
21
22
23def generate_random_input(shape, dtype):
24    return np.random.randn(*shape).astype(dtype)
25
26
27def generate_expect_forward_output(x):
28    return np.any(x)
29
30
31
32
33@test_utils.run_with_cell
34def any_forward_func(x):
35    return ms.ops.any(x)
36
37
38
39@test_utils.run_with_cell
40def any_vmap_func(x, in_axes=0):
41    return ops.vmap(any_forward_func, in_axes, out_axes=0)(x)
42
43
44@pytest.mark.level0
45@pytest.mark.env_onecard
46@pytest.mark.platform_arm_ascend_training
47@pytest.mark.platform_x86_ascend_training
48@pytest.mark.parametrize('context_mode', [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
49def test_ops_any_forward(context_mode):
50    """
51    Feature: pyboost function.
52    Description: test function any forward.
53    Expectation: expect correct result.
54    """
55    ms.context.set_context(mode=context_mode)
56    x = generate_random_input((2, 3, 4, 5), np.float32)
57    output = any_forward_func(ms.Tensor(x))
58    expect = generate_expect_forward_output(x)
59    np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
60
61
62
63@pytest.mark.level1
64@pytest.mark.env_onecard
65@pytest.mark.platform_arm_ascend_training
66@pytest.mark.platform_x86_ascend_training
67@pytest.mark.parametrize('context_mode', [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
68def test_ops_any_vmap(context_mode):
69    """
70    Feature: pyboost function.
71    Description: test function any vmap feature.
72    Expectation: expect correct result.
73    """
74    ms.context.set_context(mode=context_mode)
75    x = generate_random_input((2, 3, 4, 5), np.float32)
76    output = any_vmap_func(ms.Tensor(x), 0)
77    expect = generate_expect_forward_output(x)
78    np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
79
80
81@pytest.mark.level1
82@pytest.mark.env_onecard
83@pytest.mark.platform_arm_ascend_training
84@pytest.mark.platform_x86_ascend_training
85@pytest.mark.parametrize('context_mode', [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
86def test_ops_any_forward_dynamic_shape(context_mode):
87    """
88    Feature: pyboost function.
89    Description: test function any forward with dynamic shape.
90    Expectation: expect correct result.
91    """
92    ms.context.set_context(mode=context_mode)
93
94    x_dyn = ms.Tensor(shape=[None, None, None, None], dtype=ms.float32)
95    test_cell = test_utils.to_cell_obj(any_forward_func)
96    test_cell.set_inputs(x_dyn)
97
98    x1 = generate_random_input((2, 3, 4, 5), np.float32)
99    output = test_cell(ms.Tensor(x1))
100    expect = generate_expect_forward_output(x1)
101    np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
102
103    x2 = generate_random_input((3, 4, 5, 6), np.float32)
104    output = test_cell(ms.Tensor(x2))
105    expect = generate_expect_forward_output(x2)
106    np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
107
108
109@pytest.mark.level1
110@pytest.mark.env_onecard
111@pytest.mark.platform_arm_ascend_training
112@pytest.mark.platform_x86_ascend_training
113@pytest.mark.parametrize('context_mode', [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
114def test_ops_any_forward_dynamic_rank(context_mode):
115    """
116    Feature: pyboost function.
117    Description: test function any forward with dynamic rank.
118    Expectation: expect correct result.
119    """
120    ms.context.set_context(mode=context_mode)
121
122    x_dyn = ms.Tensor(shape=None, dtype=ms.float32)
123    test_cell = test_utils.to_cell_obj(any_forward_func)
124    test_cell.set_inputs(x_dyn)
125
126    x1 = generate_random_input((2, 3, 4, 5), np.float32)
127    output = test_cell(ms.Tensor(x1))
128    expect = generate_expect_forward_output(x1)
129    np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
130
131    x2 = generate_random_input((3, 4, 5, 6), np.float32)
132    output = test_cell(ms.Tensor(x2))
133    expect = generate_expect_forward_output(x2)
134    np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
135
136
137
138
139@pytest.mark.level1
140@pytest.mark.env_onecard
141@pytest.mark.platform_arm_ascend_training
142@pytest.mark.platform_x86_ascend_training
143@pytest.mark.parametrize('mode', ['pynative', 'KBK', 'GE'])
144def test_any_forward_static_shape(mode):
145    """
146    Feature: Test any with static shape in graph and pynative mode.
147    Description: call ops.any with valid input and index.
148    Expectation: return the correct value.
149    """
150    x = generate_random_input((3, 4, 5), np.float32)
151
152    if mode == 'pynative':
153        output = any_forward_func(ms.Tensor(x))
154    elif mode == 'KBK':
155        output = (jit(any_forward_func, jit_config=JitConfig(jit_level="O0")))(ms.Tensor(x))
156    else:
157        output = (jit(any_forward_func, jit_config=JitConfig(jit_level="O2")))(ms.Tensor(x))
158
159    expect = generate_expect_forward_output(x)
160    assert np.allclose(output.asnumpy(), expect, rtol=1e-4)
161
162
163
164@pytest.mark.level1
165@pytest.mark.env_onecard
166@pytest.mark.platform_arm_ascend_training
167@pytest.mark.platform_x86_ascend_training
168def test_any_dynamic_shape_testop():
169    """
170    Feature: Test any with dynamic shape in graph mode using TEST_OP.
171    Description: call ops.any with valid input and index.
172    Expectation: return the correct value.
173    """
174    x1 = generate_random_input((3, 4, 5), np.float32)
175    x2 = generate_random_input((3, 7, 8, 3), np.float32)
176
177    TEST_OP(any_forward_func, [[ms.Tensor(x1)], [ms.Tensor(x2)]], '', disable_yaml_check=True)
178
179
180@pytest.mark.level1
181@pytest.mark.env_onecard
182@pytest.mark.platform_arm_ascend_training
183@pytest.mark.platform_x86_ascend_training
184@pytest.mark.parametrize('param_jit_level', ["O2", "O0"])
185def test_any_vmap(param_jit_level):
186    """
187    Feature: Test any with vmap.
188    Description: call ops.any with valid input and index.
189    Expectation: return the correct value.
190    """
191    def _foreach_run(inputs, batch):
192        out = []
193        for i in range(inputs.shape[batch]):
194            if batch == -1:
195                input_inner = inputs[..., i]
196            else:
197                input_inner = inputs[i, ...]
198            out.append(any_forward_func(input_inner))
199        out = ops.Stack()(out)
200        return out
201
202    ms.set_context(jit_level=param_jit_level)
203    x = generate_random_input((4, 5, 6), np.float32)
204
205    batch_axis = -1
206    output = any_vmap_func(ms.Tensor(x), batch_axis)
207    expect = _foreach_run(ms.Tensor(x), batch_axis)
208    assert np.allclose(output.asnumpy(), expect.asnumpy(), rtol=1e-4)
209
210    batch_axis = 0
211    output = any_vmap_func(ms.Tensor(x), batch_axis)
212    expect = _foreach_run(ms.Tensor(x), batch_axis)
213    assert np.allclose(output.asnumpy(), expect.asnumpy(), rtol=1e-4)
214