Star1316
1432
代码
代码
Issues128
Pull Requests259
Actions
项目讨论
Wiki
分析
Star1316
1432
已合并
arg_max_v2，arg_max_with_value，arg_min，arg_min_with_value添加蓝区st #3362
huangzhiyuan创建于 6月16日
arg_max_v2，arg_max_with_value，arg_min，arg_min_with_value添加蓝区st #3362
已合并
huangzhiyuan创建于 6月16日
从h1234515/ops-mathgolden合入到cann/ops-mathmaster
共 8 个文件变更+178-0
Amath/arg_max_v2/tests/assets/golden.py+38-0
@@ -0,0 +1,38 @@
1+#!/usr/bin/env python3
2+# -*- coding: UTF-8 -*-
3+# ----------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# ----------------------------------------------------------------------------
12+import numpy as np
13+ 
14+__golden__ = {
15+ "kernel": {
16+ "arg_max_v2": "arg_max_v2_golden"
17+ }
18+}
19+ 
20+ 
21+def arg_max_v2_golden(x, dimension, *, dtype=None, **kwargs):
22+ '''
23+ Kernel golden for arg_max_v2.
24+ All the parameters follow @arg_max_v2_def.cpp without outputs.
25+ All the input Tensors are numpy.ndarray.
26+ kwargs may contain: short_soc_version, input_ori_shapes, output_ori_shapes,
27+ input_formats, output_formats, input_ori_formats, output_ori_formats,
28+ input_dtypes, output_dtypes.
29+ '''
30+ ori_dtype = kwargs.get("input_dtypes", ["float32"])[0]
31+ output_dtype = kwargs.get("output_dtypes", ["int32"])[0]
32+ 
33+ if "bfloat16" in str(ori_dtype).lower():
34+ x = x.astype("float32")
35+ 
36+ axis = int(dimension) % len(x.shape)
37+ indices = np.argmax(x, axis=axis).astype(output_dtype, copy=False)
38+ return indices
Amath/arg_max_v2/tests/st/ttk_kernel_arg_max_v2_st.csv+5-0
@@ -0,0 +1,5 @@
1+testcase_name,op_name,input_shapes,input_ori_shapes,output_shapes,output_ori_shapes,input_dtypes,output_dtypes,attributes,output_formats,output_ori_formats,input_formats,input_ori_formats,stc_op_name,const_input_indexes,input_as_list_distribution,input_data_ranges,shape_check,absolute_precision,dump_input_data_name,dump_output_data_name,precision_tolerances,strict_precision_mode,is_enabled,priority
2+ArgMaxV2_AR_cutR_fp16,arg_max_v2,"((64, 60000), ())","((64, 60000), ())","((64,),)","((64,),)","('float16', 'int64')","('int64',)","{'dimension': -1, 'dtype': 9}","('ND',)","('ND',)","('ND', 'ND')","('ND', 'ND')",,"(1,)",,"((-inf, inf),)",,1e-08,,,,,1,0
3+ArgMaxV2_ARA_cutA_fp32,arg_max_v2,"((64, 1024, 16), ())","((64, 1024, 16), ())","((64, 16),)","((64, 16),)","('float32', 'int64')","('int64',)","{'dimension': 1, 'dtype': 9}","('ND',)","('ND',)","('ND', 'ND')","('ND', 'ND')",,"(1,)",,"((-inf, inf),)",,1e-08,,,,,1,0
4+ArgMaxV2_RA_cutR_bf16,arg_max_v2,"((64, 61440), ())","((64, 61440), ())","((61440,),)","((61440,),)","('bfloat16', 'int64')","('int64',)","{'dimension': 0, 'dtype': 9}","('ND',)","('ND',)","('ND', 'ND')","('ND', 'ND')",,"(1,)",,"((-inf, inf),)",,1e-08,,,,,1,0
5+ArgMaxV2_ARGather_fp16,arg_max_v2,"((8192, 64), ())","((8192, 64), ())","((8192,),)","((8192,),)","('float16', 'int64')","('int64',)","{'dimension': -1, 'dtype': 9}","('ND',)","('ND',)","('ND', 'ND')","('ND', 'ND')",,"(1,)",,"((-inf, inf),)",,1e-08,,,,,1,0
Amath/arg_max_with_value/tests/assets/golden.py+41-0
@@ -0,0 +1,41 @@
1+#!/usr/bin/env python3
2+# -*- coding: UTF-8 -*-
3+# ----------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# ----------------------------------------------------------------------------
12+import numpy as np
13+ 
14+__golden__ = {
15+ "kernel": {
16+ "arg_max_with_value": "arg_max_with_value_golden"
17+ }
18+}
19+ 
20+ 
21+def arg_max_with_value_golden(x, *, dimension, keep_dims=False, indice_dtype=None, **kwargs):
22+ '''
23+ Kernel golden for arg_max_with_value.
24+ All the parameters follow @arg_max_with_value_def.cpp without outputs.
25+ All the input Tensors are numpy.ndarray.
26+ kwargs may contain: short_soc_version, input_ori_shapes, output_ori_shapes,
27+ input_formats, output_formats, input_ori_formats, output_ori_formats,
28+ input_dtypes, output_dtypes.
29+ '''
30+ ori_dtype = kwargs.get("input_dtypes", ["float32"])[0]
31+ output_dtypes = kwargs.get("output_dtypes", ["int32", "float32"])
32+ indice_dtype = output_dtypes[0] if indice_dtype is None else indice_dtype
33+ 
34+ x_bak = x
35+ if "bfloat16" in str(ori_dtype).lower():
36+ x = x.astype("float32")
37+ 
38+ axis = dimension % len(x.shape)
39+ indices = np.argmax(x, axis=axis).astype(indice_dtype, copy=False)
40+ values = np.take_along_axis(x_bak, np.expand_dims(indices, axis=axis), axis=axis)
41+ return indices, values
Amath/arg_max_with_value/tests/st/ttk_kernel_arg_max_with_value_st.csv+5-0
@@ -0,0 +1,5 @@
1+testcase_name,op_name,input_shapes,input_ori_shapes,output_shapes,output_ori_shapes,input_dtypes,output_dtypes,attributes,output_formats,output_ori_formats,input_formats,input_ori_formats,stc_op_name,const_input_indexes,input_as_list_distribution,input_data_ranges,shape_check,absolute_precision,dump_input_data_name,dump_output_data_name,precision_tolerances,strict_precision_mode,is_enabled,priority
2+ArgMaxWithValue_AR_cutR_fp16,arg_max_with_value,"((64, 60000),)","((64, 60000),)","((64,), (64,))","((64,), (64,))","('float16',)","('int32', 'float16')",{'dimension': -1},"('ND', 'ND')","('ND', 'ND')","('ND',)","('ND',)",,,,"((-inf, inf),)",,1e-08,,,,,1,0
3+ArgMaxWithValue_ARA_cutA_fp32,arg_max_with_value,"((64, 1024, 16),)","((64, 1024, 16),)","((64, 16), (64, 16))","((64, 16), (64, 16))","('float32',)","('int32', 'float32')",{'dimension': 1},"('ND', 'ND')","('ND', 'ND')","('ND',)","('ND',)",,,,"((-inf, inf),)",,1e-08,,,,,1,0
4+ArgMaxWithValue_RA_cutR_bf16,arg_max_with_value,"((64, 61440),)","((64, 61440),)","((61440,), (61440,))","((61440,), (61440,))","('bfloat16',)","('int32', 'bfloat16')",{'dimension': 0},"('ND', 'ND')","('ND', 'ND')","('ND',)","('ND',)",,,,"((-inf, inf),)",,1e-08,,,,,1,0
5+ArgMaxWithValue_ARGather_fp16,arg_max_with_value,"((8192, 64),)","((8192, 64),)","((8192,), (8192,))","((8192,), (8192,))","('float16',)","('int32', 'float16')",{'dimension': -1},"('ND', 'ND')","('ND', 'ND')","('ND',)","('ND',)",,,,"((-inf, inf),)",,1e-08,,,,,1,0
Amath/arg_min/tests/assets/golden.py+38-0
@@ -0,0 +1,38 @@
1+#!/usr/bin/env python3
2+# -*- coding: UTF-8 -*-
3+# ----------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# ----------------------------------------------------------------------------
12+import numpy as np
13+ 
14+__golden__ = {
15+ "kernel": {
16+ "arg_min": "arg_min_golden"
17+ }
18+}
19+ 
20+ 
21+def arg_min_golden(x, dimension, *, dtype=None, **kwargs):
22+ '''
23+ Kernel golden for arg_min.
24+ All the parameters follow @arg_min_def.cpp without outputs.
25+ All the input Tensors are numpy.ndarray.
26+ kwargs may contain: short_soc_version, input_ori_shapes, output_ori_shapes,
27+ input_formats, output_formats, input_ori_formats, output_ori_formats,
28+ input_dtypes, output_dtypes.
29+ '''
30+ ori_dtype = kwargs.get("input_dtypes", ["float32"])[0]
31+ output_dtype = kwargs.get("output_dtypes", ["int32"])[0]
32+ 
33+ if "bfloat16" in str(ori_dtype).lower():
34+ x = x.astype("float32")
35+ 
36+ axis = int(dimension) % len(x.shape)
37+ indices = np.argmin(x, axis=axis).astype(output_dtype, copy=False)
38+ return indices
Amath/arg_min/tests/st/ttk_kernel_arg_min_st.csv+5-0
@@ -0,0 +1,5 @@
1+testcase_name,op_name,input_shapes,input_ori_shapes,output_shapes,output_ori_shapes,input_dtypes,output_dtypes,attributes,output_formats,output_ori_formats,input_formats,input_ori_formats,stc_op_name,const_input_indexes,input_as_list_distribution,input_data_ranges,shape_check,absolute_precision,dump_input_data_name,dump_output_data_name,precision_tolerances,strict_precision_mode,is_enabled,priority
2+ArgMin_AR_cutR_fp16,arg_min,"((64, 60000), ())","((64, 60000), ())","((64,),)","((64,),)","('float16', 'int32')","('int32',)",{'dimension': -1},"('ND',)","('ND',)","('ND', 'ND')","('ND', 'ND')",,"(1,)",,"((-inf, inf),)",,1e-08,,,,,1,0
3+ArgMin_ARA_cutA_fp32,arg_min,"((64, 1024, 16), ())","((64, 1024, 16), ())","((64, 16),)","((64, 16),)","('float32', 'int32')","('int32',)",{'dimension': 1},"('ND',)","('ND',)","('ND', 'ND')","('ND', 'ND')",,"(1,)",,"((-inf, inf),)",,1e-08,,,,,1,0
4+ArgMin_RA_cutR_bf16,arg_min,"((64, 61440), ())","((64, 61440), ())","((61440,),)","((61440,),)","('bfloat16', 'int32')","('int32',)",{'dimension': 0},"('ND',)","('ND',)","('ND', 'ND')","('ND', 'ND')",,"(1,)",,"((-inf, inf),)",,1e-08,,,,,1,0
5+ArgMin_ARGather_fp16,arg_min,"((8192, 64), ())","((8192, 64), ())","((8192,),)","((8192,),)","('float16', 'int32')","('int32',)",{'dimension': -1},"('ND',)","('ND',)","('ND', 'ND')","('ND', 'ND')",,"(1,)",,"((-inf, inf),)",,1e-08,,,,,1,0
Amath/arg_min_with_value/tests/assets/golden.py+41-0
@@ -0,0 +1,41 @@
1+#!/usr/bin/env python3
2+# -*- coding: UTF-8 -*-
3+# ----------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# ----------------------------------------------------------------------------
12+import numpy as np
13+ 
14+__golden__ = {
15+ "kernel": {
16+ "arg_min_with_value": "arg_min_with_value_golden"
17+ }
18+}
19+ 
20+ 
21+def arg_min_with_value_golden(x, *, dimension, keep_dims=False, indice_dtype=None, **kwargs):
22+ '''
23+ Kernel golden for arg_min_with_value.
24+ All the parameters follow @arg_min_with_value_def.cpp without outputs.
25+ All the input Tensors are numpy.ndarray.
26+ kwargs may contain: short_soc_version, input_ori_shapes, output_ori_shapes,
27+ input_formats, output_formats, input_ori_formats, output_ori_formats,
28+ input_dtypes, output_dtypes.
29+ '''
30+ ori_dtype = kwargs.get("input_dtypes", ["float32"])[0]
31+ output_dtypes = kwargs.get("output_dtypes", ["int32", "float32"])
32+ indice_dtype = output_dtypes[0] if indice_dtype is None else indice_dtype
33+ 
34+ x_bak = x
35+ if "bfloat16" in str(ori_dtype).lower():
36+ x = x.astype("float32")
37+ 
38+ axis = dimension % len(x.shape)
39+ indices = np.argmin(x, axis=axis).astype(indice_dtype, copy=False)
40+ values = np.take_along_axis(x_bak, np.expand_dims(indices, axis=axis), axis=axis)
41+ return indices, values
Amath/arg_min_with_value/tests/st/ttk_kernel_arg_min_with_value_st.csv+5-0
@@ -0,0 +1,5 @@
1+testcase_name,op_name,input_shapes,input_ori_shapes,output_shapes,output_ori_shapes,input_dtypes,output_dtypes,attributes,output_formats,output_ori_formats,input_formats,input_ori_formats,stc_op_name,const_input_indexes,input_as_list_distribution,input_data_ranges,shape_check,absolute_precision,dump_input_data_name,dump_output_data_name,precision_tolerances,strict_precision_mode,is_enabled,priority
2+ArgMinWithValue_AR_cutR_fp16,arg_min_with_value,"((64, 60000),)","((64, 60000),)","((64,), (64,))","((64,), (64,))","('float16',)","('int32', 'float16')",{'dimension': -1},"('ND', 'ND')","('ND', 'ND')","('ND',)","('ND',)",,,,"((-inf, inf),)",,1e-08,,,,,1,0
3+ArgMinWithValue_ARA_cutA_fp32,arg_min_with_value,"((64, 1024, 16),)","((64, 1024, 16),)","((64, 16), (64, 16))","((64, 16), (64, 16))","('float32',)","('int32', 'float32')",{'dimension': 1},"('ND', 'ND')","('ND', 'ND')","('ND',)","('ND',)",,,,"((-inf, inf),)",,1e-08,,,,,1,0
4+ArgMinWithValue_RA_cutR_bf16,arg_min_with_value,"((64, 61440),)","((64, 61440),)","((61440,), (61440,))","((61440,), (61440,))","('bfloat16',)","('int32', 'bfloat16')",{'dimension': 0},"('ND', 'ND')","('ND', 'ND')","('ND',)","('ND',)",,,,"((-inf, inf),)",,1e-08,,,,,1,0
5+ArgMinWithValue_ARGather_fp16,arg_min_with_value,"((8192, 64),)","((8192, 64),)","((8192,), (8192,))","((8192,), (8192,))","('float16',)","('int32', 'float16')",{'dimension': -1},"('ND', 'ND')","('ND', 'ND')","('ND',)","('ND',)",,,,"((-inf, inf),)",,1e-08,,,,,1,0