mindspore 2.0.0a0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops/_op_impl/tbe/batch_to_space.py

@@ -16,7 +16,7 @@
 """BatchToSpace op"""
 from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
 
-batch_to_space_op_info = TBERegOp("BatchToSpace") \
+batch_to_space_op_info = TBERegOp("BatchToSpaceD") \
     .fusion_type("OPAQUE") \
     .async_flag(False) \
     .binfile_name("batch_to_space_d.so") \

mindspore/ops/_op_impl/tbe/batch_to_space_nd.py

@@ -16,7 +16,7 @@
 """BatchToSpaceND op"""
 from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
 
-batch_to_space_nd_op_info = TBERegOp("BatchToSpaceND") \
+batch_to_space_nd_op_info = TBERegOp("BatchToSpaceNDD") \
     .fusion_type("OPAQUE") \
     .async_flag(False) \
     .binfile_name("batch_to_space_nd_d.so") \

mindspore/ops/_op_impl/tbe/bn_infer_grad.py

@@ -23,11 +23,13 @@ bn_infer_grad_op_info = TBERegOp("BNInferGrad") \
     .compute_cost(10) \
     .kernel_name("bn_infer_grad") \
     .partial_flag(True) \
+    .dynamic_shape(True)\
+    .dynamic_compile_static(True)\
     .attr("epsilon", "optional", "float", "all", "0.0001") \
-    .input(0, "grads", False, "required", "all", reshape_type="NC") \
+    .input(0, "grads", False, "required", "all", reshape_type="NCH") \
     .input(1, "scale", False, "required", "all") \
     .input(2, "batch_variance", False, "required", "all") \
-    .output(0, "x_backprop", False, "required", "all", reshape_type="NC") \
+    .output(0, "x_backprop", False, "required", "all", reshape_type="NCH") \
     .dtype_format(DataType.F16_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F16_5HD) \
     .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \
     .get_op_info()

mindspore/ops/_op_impl/tbe/bn_training_update.py

@@ -25,7 +25,6 @@ bn_training_update_op_info = TBERegOp("BNTrainingUpdate") \
     .partial_flag(True) \
     .attr("factor", "optional", "float", "all") \
     .attr("epsilon", "optional", "float", "all") \
-    .attr("isRef", "optional", "bool", "all", "true") \
     .input(0, "x", False, "required", "all", reshape_type="NC") \
     .input(1, "sum", False, "required", "all") \
     .input(2, "square_sum", False, "required", "all") \

mindspore/ops/_op_impl/tbe/bn_training_update_ds.py

@@ -26,7 +26,6 @@ bn_training_update_op_info = TBERegOp("BNTrainingUpdate") \
     .dynamic_shape(True) \
     .attr("factor", "optional", "float", "all") \
     .attr("epsilon", "optional", "float", "all") \
-    .attr("isRef", "optional", "bool", "all", "true") \
     .input(0, "x", False, "required", "all", reshape_type="NC") \
     .input(1, "sum", False, "required", "all") \
     .input(2, "square_sum", False, "required", "all") \

mindspore/ops/_op_impl/tbe/broadcast_to_ds.py

@@ -27,12 +27,14 @@ broadcast_to_op_info = TBERegOp("DynamicBroadcastTo") \
     .input(0, "x", False, "required", "all") \
     .input(1, "shape", False, "required", "all") \
     .output(0, "y", False, "required", "all") \
-    .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.F16_Default) \
-    .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.F32_Default) \
     .dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
-    .dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.F16_Default) \
-    .dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.F32_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.F32_Default) \
+    .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.F16_Default) \
     .dtype_format(DataType.I32_Default, DataType.I64_Default, DataType.I32_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.F32_Default) \
+    .dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.F16_Default) \
+    .dtype_format(DataType.U32_Default, DataType.I32_Default, DataType.U32_Default) \
+    .dtype_format(DataType.U32_Default, DataType.I64_Default, DataType.U32_Default) \
     .get_op_info()
 
 

mindspore/ops/_op_impl/tbe/cast.py

@@ -41,8 +41,6 @@ cast_op_info = TBERegOp("Cast") \
     .dtype_format(DataType.I32_None, DataType.BOOL_None) \
     .dtype_format(DataType.I32_None, DataType.F16_None) \
     .dtype_format(DataType.I32_None, DataType.F32_None) \
-    .dtype_format(DataType.I32_None, DataType.I8_None) \
-    .dtype_format(DataType.I32_None, DataType.U8_None) \
     .dtype_format(DataType.F16_None, DataType.U8_None) \
     .dtype_format(DataType.F16_None, DataType.F32_None) \
     .dtype_format(DataType.F16_None, DataType.I32_None) \

mindspore/ops/_op_impl/tbe/cast_ds.py

@@ -25,7 +25,8 @@ cast_ds_op_info = TBERegOp("Cast") \
     .partial_flag(True) \
     .need_check_supported(True) \
     .attr("dst_type", "required", "int", "all") \
-    .dynamic_shape(True)\
+    .dynamic_shape(True) \
+    .dynamic_compile_static(True) \
     .input(0, "x", False, "required", "all") \
     .output(0, "y", False, "required", "all") \
     .op_pattern("formatAgnostic") \
@@ -42,11 +43,10 @@ cast_ds_op_info = TBERegOp("Cast") \
     .dtype_format(DataType.I32_None, DataType.BOOL_None) \
     .dtype_format(DataType.I32_None, DataType.F16_None) \
     .dtype_format(DataType.I32_None, DataType.F32_None) \
-    .dtype_format(DataType.I32_None, DataType.I8_None) \
-    .dtype_format(DataType.I32_None, DataType.U8_None) \
     .dtype_format(DataType.F16_None, DataType.U8_None) \
     .dtype_format(DataType.F16_None, DataType.F32_None) \
     .dtype_format(DataType.F16_None, DataType.I32_None) \
+    .dtype_format(DataType.F16_None, DataType.BOOL_None) \
     .dtype_format(DataType.F32_None, DataType.F16_None) \
     .dtype_format(DataType.F32_None, DataType.I32_None) \
     .get_op_info()

mindspore/ops/_op_impl/tbe/data_format_dim_map_ds.py

@@ -24,6 +24,7 @@ data_format_dim_map_op_info = TBERegOp("DataFormatDimMap") \
     .kernel_name("data_format_dim_map") \
     .partial_flag(True) \
     .dynamic_shape(True) \
+    .dynamic_compile_static(True) \
     .attr("dst_format", "optional", "str", "all") \
     .attr("src_format", "optional", "str", "all") \
     .input(0, "x", False, "required", "all") \

mindspore/ops/_op_impl/tbe/depthwise_conv2d.py

@@ -16,7 +16,7 @@
 """DepthwiseConv2D op"""
 from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
 
-depthwise_conv2d_op_info = TBERegOp("DepthwiseConv2dNative") \
+depthwise_conv2d_op_info = TBERegOp("DepthwiseConv2D") \
     .fusion_type("CONVOLUTION") \
     .async_flag(False) \
     .binfile_name("depthwise_conv2d.so") \
@@ -33,7 +33,7 @@ depthwise_conv2d_op_info = TBERegOp("DepthwiseConv2dNative") \
     .input(2, "bias", False, "optional", "all") \
     .input(3, "offset_w", False, "optional", "all") \
     .output(0, "y", True, "required", "all") \
-    .dtype_format(DataType.F16_5HD, DataType.F16_FracZ, DataType.F16_Default, DataType.F16_Default,
+    .dtype_format(DataType.F16_5HD, DataType.F16_C1HWNCoC0, DataType.F16_Default, DataType.F16_Default,
                   DataType.F16_5HD) \
     .get_op_info()
 

mindspore/ops/_op_impl/tbe/dynamic_atomic_addr_clean.py

@@ -24,7 +24,7 @@ dynamic_atomic_addr_clean_op_info = TBERegOp("DynamicAtomicAddrClean") \
     .kernel_name("dynamic_atomic_addr_clean") \
     .partial_flag(True) \
     .dynamic_shape(True) \
-    .attr("automic_add_mem_size", "required", "listUInt64", "all") \
+    .attr("automic_add_mem_size", "required", "listInt", "all") \
     .get_op_info()
 
 

mindspore/ops/_op_impl/tbe/gather_nd.py

@@ -23,6 +23,7 @@ gather_nd_op_info = TBERegOp("GatherNd") \
     .compute_cost(10) \
     .kernel_name("gather_nd") \
     .partial_flag(True) \
+    .attr("negative_index_support", "optional", "bool", "all", "false") \
     .dynamic_compile_static(True) \
     .dynamic_shape(True) \
     .input(0, "x1", False, "required", "all") \

mindspore/ops/_op_impl/tbe/{index_add.py → inplace_index_add.py}

@@ -16,7 +16,7 @@
 """IndexAdd op"""
 from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
 
-index_add_op_info = TBERegOp("IndexAdd") \
+inplace_index_add_op_info = TBERegOp("InplaceIndexAdd") \
     .fusion_type("OPAQUE") \
     .async_flag(False) \
     .binfile_name("inplace_index_add.so") \
@@ -32,14 +32,11 @@ index_add_op_info = TBERegOp("IndexAdd") \
     .output(0, "input_x", False, "required", "all") \
     .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.F16_Default, DataType.F16_Default) \
     .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.F32_Default, DataType.F32_Default) \
-    .dtype_format(DataType.U8_Default, DataType.I32_Default, DataType.U8_Default, DataType.U8_Default) \
-    .dtype_format(DataType.I8_Default, DataType.I32_Default, DataType.I8_Default, DataType.I8_Default) \
-    .dtype_format(DataType.I16_Default, DataType.I32_Default, DataType.I16_Default, DataType.I16_Default) \
     .dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
     .get_op_info()
 
 
-@op_info_register(index_add_op_info)
-def _index_add_tbe():
+@op_info_register(inplace_index_add_op_info)
+def _inplace_index_add_tbe():
     """IndexAdd TBE register"""
     return

mindspore/ops/_op_impl/tbe/matmul_ds.py

@@ -36,6 +36,8 @@ matmul_op_info = TBERegOp("MatMul") \
                   DataType.F16_FracNZ) \
     .dtype_format(DataType.F16_FracNZ, DataType.F16_FracNZ, DataType.F32_Default, DataType.I8_Default,
                   DataType.F16_FracNZ) \
+    .dtype_format(DataType.F16_FracNZ, DataType.F16_FracNZ, DataType.F32_Default, DataType.I8_Default,
+                  DataType.F32_FracNZ) \
     .get_op_info()
 
 

mindspore/ops/_op_impl/tbe/npu_clear_float_status_v2.py

@@ -0,0 +1,35 @@
+# Copyright 2023 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.
+# ============================================================================
+
+"""NPUClearFloatStatusV2 op"""
+from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
+
+npu_clear_float_status_v2_op_info = TBERegOp("NPUClearFloatStatusV2") \
+    .fusion_type("OPAQUE") \
+    .async_flag(False) \
+    .binfile_name("n_p_u_clear_float_status_v2.so") \
+    .compute_cost(10) \
+    .kernel_name("n_p_u_clear_float_status_v2") \
+    .partial_flag(True) \
+    .input(0, "addr", False, "required", "all") \
+    .output(0, "data", False, "required", "all") \
+    .dtype_format(DataType.I32_Default, DataType.I32_Default) \
+    .get_op_info()
+
+
+@op_info_register(npu_clear_float_status_v2_op_info)
+def _npu_clear_float_status_v2_tbe():
+    """NPUClearFloatStatusV2 TBE register"""
+    return

mindspore/ops/_op_impl/tbe/npu_get_float_status_v2.py

@@ -0,0 +1,35 @@
+# Copyright 2023 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.
+# ============================================================================
+
+"""NPUGetFloatStatusV2 op"""
+from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
+
+npu_get_float_status_v2_op_info = TBERegOp("NPUGetFloatStatusV2") \
+    .fusion_type("ELEMWISE") \
+    .async_flag(False) \
+    .binfile_name("n_p_u_get_float_status_v2.so") \
+    .compute_cost(10) \
+    .kernel_name("n_p_u_get_float_status_v2") \
+    .partial_flag(True) \
+    .input(0, "addr", False, "required", "all") \
+    .output(0, "data", False, "required", "all") \
+    .dtype_format(DataType.I32_Default, DataType.I32_Default) \
+    .get_op_info()
+
+
+@op_info_register(npu_get_float_status_v2_op_info)
+def _npu_get_float_status_v2_tbe():
+    """NPUGetFloatStatusV2 TBE register"""
+    return

mindspore/ops/_op_impl/tbe/scatter_mul.py

@@ -33,6 +33,8 @@ scatter_mul_op_info = TBERegOp("ScatterMul") \
     .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.F16_Default, DataType.F16_Default) \
     .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.F32_Default, DataType.F32_Default) \
     .dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
+    .dtype_format(DataType.I8_Default, DataType.I32_Default, DataType.I8_Default, DataType.I8_Default) \
+    .dtype_format(DataType.U8_Default, DataType.I32_Default, DataType.U8_Default, DataType.U8_Default) \
     .get_op_info()
 
 

mindspore/ops/_op_impl/tbe/scatter_nd_add.py

@@ -32,8 +32,6 @@ scatter_nd_add_op_info = TBERegOp("ScatterNdAdd") \
     .output(0, "var", False, "required", "all") \
     .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.F16_Default, DataType.F16_Default) \
     .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.F32_Default, DataType.F32_Default) \
-    .dtype_format(DataType.I8_Default, DataType.I32_Default, DataType.I8_Default, DataType.I8_Default) \
-    .dtype_format(DataType.U8_Default, DataType.I32_Default, DataType.U8_Default, DataType.U8_Default) \
     .dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
     .get_op_info()
 

mindspore/ops/_op_impl/tbe/space_to_batch.py

@@ -16,7 +16,7 @@
 """SpaceToBatch op"""
 from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
 
-space_to_batch_op_info = TBERegOp("SpaceToBatch") \
+space_to_batch_op_info = TBERegOp("SpaceToBatchD") \
     .fusion_type("OPAQUE") \
     .async_flag(False) \
     .binfile_name("space_to_batch_d.so") \

mindspore/ops/_op_impl/tbe/space_to_batch_nd.py

@@ -16,7 +16,7 @@
 """SpaceToBatchND op"""
 from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
 
-space_to_batch_nd_op_info = TBERegOp("SpaceToBatchND") \
+space_to_batch_nd_op_info = TBERegOp("SpaceToBatchNDD") \
     .fusion_type("OPAQUE") \
     .async_flag(False) \
     .binfile_name("space_to_batch_nd_d.so") \

mindspore/ops/_op_impl/tbe/trans_data_ds.py

@@ -59,10 +59,6 @@ trans_data_op_info = TBERegOp("TransData") \
     .dtype_format(DataType.F32_Default, DataType.F32_FracNZ) \
     .dtype_format(DataType.F16_FracNZ, DataType.F16_Default) \
     .dtype_format(DataType.F32_FracNZ, DataType.F32_Default) \
-    .dtype_format(DataType.BOOL_NHWC, DataType.BOOL_5HD) \
-    .dtype_format(DataType.BOOL_Default, DataType.BOOL_5HD) \
-    .dtype_format(DataType.BOOL_5HD, DataType.BOOL_NHWC) \
-    .dtype_format(DataType.BOOL_5HD, DataType.BOOL_Default) \
     .dtype_format(DataType.F16_Default, DataType.F16_NHWC) \
     .dtype_format(DataType.F16_Default, DataType.F16_HWCN) \
     .dtype_format(DataType.F16_NHWC, DataType.F16_Default) \
@@ -83,6 +79,10 @@ trans_data_op_info = TBERegOp("TransData") \
     .dtype_format(DataType.I8_NHWC, DataType.I8_HWCN) \
     .dtype_format(DataType.I8_HWCN, DataType.I8_Default) \
     .dtype_format(DataType.I8_HWCN, DataType.I8_NHWC) \
+    .dtype_format(DataType.I8_Default, DataType.I8_NC1HWC0) \
+    .dtype_format(DataType.I8_HWCN, DataType.I8_C1HWNCoC0) \
+    .dtype_format(DataType.I8_NCDHW, DataType.I8_NDC1HWC0) \
+    .dtype_format(DataType.I8_NDC1HWC0, DataType.I8_NCDHW) \
     .dtype_format(DataType.I16_Default, DataType.I16_NHWC) \
     .dtype_format(DataType.I16_Default, DataType.I16_HWCN) \
     .dtype_format(DataType.I16_NHWC, DataType.I16_Default) \
@@ -95,6 +95,8 @@ trans_data_op_info = TBERegOp("TransData") \
     .dtype_format(DataType.I32_NHWC, DataType.I32_HWCN) \
     .dtype_format(DataType.I32_HWCN, DataType.I32_Default) \
     .dtype_format(DataType.I32_HWCN, DataType.I32_NHWC) \
+    .dtype_format(DataType.I32_NDC1HWC0, DataType.I32_NCDHW) \
+    .dtype_format(DataType.I32_NCDHW, DataType.I32_NDC1HWC0) \
     .dtype_format(DataType.I64_Default, DataType.I64_NHWC) \
     .dtype_format(DataType.I64_Default, DataType.I64_HWCN) \
     .dtype_format(DataType.I64_NHWC, DataType.I64_Default) \
@@ -107,6 +109,9 @@ trans_data_op_info = TBERegOp("TransData") \
     .dtype_format(DataType.U8_NHWC, DataType.U8_HWCN) \
     .dtype_format(DataType.U8_HWCN, DataType.U8_Default) \
     .dtype_format(DataType.U8_HWCN, DataType.U8_NHWC) \
+    .dtype_format(DataType.U8_Default, DataType.U8_NC1HWC0) \
+    .dtype_format(DataType.U8_NCDHW, DataType.U8_NDC1HWC0) \
+    .dtype_format(DataType.U8_NDC1HWC0, DataType.U8_NCDHW) \
     .dtype_format(DataType.U16_Default, DataType.U16_NHWC) \
     .dtype_format(DataType.U16_Default, DataType.U16_HWCN) \
     .dtype_format(DataType.U16_NHWC, DataType.U16_Default) \
@@ -128,7 +133,6 @@ trans_data_op_info = TBERegOp("TransData") \
     .dtype_format(DataType.I32_FracNZ, DataType.I32_Default) \
     .dtype_format(DataType.F16_NDHWC, DataType.F16_5HD) \
     .dtype_format(DataType.F16_5HD, DataType.F16_NDHWC) \
-    .dtype_format(DataType.I8_HWCN, DataType.I8_C1HWNCoC0) \
     .dtype_format(DataType.F16_HWCN, DataType.F16_FracZ) \
     .dtype_format(DataType.F16_FracZ, DataType.F16_HWCN) \
     .dtype_format(DataType.F16_HWCN, DataType.F16_FracNZ) \
@@ -159,6 +163,12 @@ trans_data_op_info = TBERegOp("TransData") \
     .dtype_format(DataType.F32_NCDHW, DataType.F32_NDC1HWC0) \
     .dtype_format(DataType.F32_NDC1HWC0, DataType.F32_NCDHW) \
     .dtype_format(DataType.F32_NCDHW, DataType.F32_NDC1HWC0) \
+    .dtype_format(DataType.I8_Default, DataType.I8_5HD) \
+    .dtype_format(DataType.I8_5HD, DataType.I8_Default) \
+    .dtype_format(DataType.U8_Default, DataType.U8_5HD) \
+    .dtype_format(DataType.U8_5HD, DataType.U8_Default) \
+    .dtype_format(DataType.I32_Default, DataType.I32_5HD) \
+    .dtype_format(DataType.I32_5HD, DataType.I32_Default) \
     .get_op_info()
 
 

mindspore/ops/_register_for_op.py

@@ -63,6 +63,7 @@ class PyFuncRegistry(UserDict):
 
 
 class OpaquePredicateRegistry(PyFuncRegistry):
+    """Registry opaque predicate functions used for dynamic obfuscation"""
     def __init__(self):
         super(OpaquePredicateRegistry, self).__init__()
         self.func_names = []

mindspore/ops/_utils/__init__.py

@@ -15,7 +15,6 @@
 
 """ops utils."""
 from __future__ import absolute_import
-
-from mindspore.ops._utils.utils import get_broadcast_shape, get_concat_offset, is_shape_unknown, is_dim_unknown
+from mindspore.ops._utils.utils import get_broadcast_shape, get_concat_offset
 
 __all__ = ['get_broadcast_shape', 'get_concat_offset']

mindspore/ops/_utils/utils.py

@@ -17,13 +17,13 @@
 from __future__ import absolute_import
 
 from mindspore.common.tensor import Tensor
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
+from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
-from mindspore.ops.primitive import constexpr
+from mindspore.ops.primitive import _primexpr
+from mindspore.common._utils import is_dim_unknown
 
 
-def get_broadcast_shape(x_shape, y_shape, prim_name, shape_type="", arg_name1="x", arg_name2="y"):
+def get_broadcast_shape(x_shape, y_shape, prim_name, arg_name1="x", arg_name2="y"):
     """
     Doing broadcast between tensor x and tensor y.
 
@@ -31,7 +31,6 @@ def get_broadcast_shape(x_shape, y_shape, prim_name, shape_type="", arg_name1="x
         x_shape (list): The shape of tensor x.
         y_shape (list): The shape of tensor y.
         prim_name (str): Primitive name.
-        shape_type (str): The type of shape, optional values are "", "min_shape" and "max_shape".
         arg_name1 (str): The arg name of x_shape.
         arg_name2 (str): The arg name of y_shape.
 
@@ -68,16 +67,11 @@ def get_broadcast_shape(x_shape, y_shape, prim_name, shape_type="", arg_name1="x
         elif x_shape[i] == -1 or y_shape[i] == -1:
             broadcast_shape_back.append(-1)
         else:
-            if shape_type == "min_shape":
-                broadcast_shape_back.append(max(x_shape[i], y_shape[i]))
-            elif shape_type == "max_shape":
-                broadcast_shape_back.append(min(x_shape[i], y_shape[i]))
-            else:
-                raise ValueError(f"For '{prim_name}', {arg_name1}.shape and {arg_name2}.shape need to "
-                                 f"broadcast. The value of {arg_name1}.shape[{i}] or {arg_name2}.shape[{i}]"
-                                 f" must be 1 or -1 when they are not the same, "
-                                 f"but got {arg_name1}.shape = {x_shape} "
-                                 f"and {arg_name2}.shape = {y_shape}.")
+            raise ValueError(f"For '{prim_name}', {arg_name1}.shape and {arg_name2}.shape need to "
+                             f"broadcast. The value of {arg_name1}.shape[{i}] or {arg_name2}.shape[{i}]"
+                             f" must be 1 or -1 when they are not the same, "
+                             f"but got {arg_name1}.shape = {x_shape} "
+                             f"and {arg_name2}.shape = {y_shape}.")
 
     broadcast_shape_front = y_shape[0: y_len - length] if length == x_len else x_shape[0: x_len - length]
     broadcast_shape = list(broadcast_shape_front) + broadcast_shape_back
@@ -92,10 +86,11 @@ def get_concat_offset(x_shp, x_type, axis, prim_name):
     validator.check_positive_int(len(x_shp[0]), "len of x_shp[0]", prim_name)
     rank_base = len(x_shp[0])
     for i in range(1, len(x_shp)):
-        validator.check('len of x_shp[%d]' % i, len(x_shp[i]), 'len of x_shp[0]', len(x_shp[0]), Rel.EQ, prim_name)
-        validator.check('x_type[%d]' % i, x_type[i], 'x_type[0]', x_type[0], Rel.EQ, prim_name)
+        validator.check('len of x_shp[%d]' % i, len(x_shp[i]), 'len of x_shp[0]',
+                        len(x_shp[0]), validator.EQ, prim_name)
+        validator.check('x_type[%d]' % i, x_type[i], 'x_type[0]', x_type[0], validator.EQ, prim_name)
 
-    validator.check_int_range(axis, -rank_base, rank_base - 1, Rel.INC_BOTH, 'axis', prim_name)
+    validator.check_int_range(axis, -rank_base, rank_base - 1, validator.INC_BOTH, 'axis', prim_name)
     if axis < 0:
         axis = axis + rank_base
     all_shp = x_shp[0][axis]
@@ -114,14 +109,14 @@ def get_concat_offset(x_shp, x_type, axis, prim_name):
     return offset, all_shp, axis
 
 
-@constexpr
+@_primexpr
 def range_op(start, limit, delta, dtype):
     """helper function to get tensor in specified range."""
     output_tensor = Tensor(list(range(start, limit, delta)), dtype)
     return output_tensor
 
 
-@constexpr
+@_primexpr
 def get_1d_shape(in_shape):
     """helper function to get 1d shape."""
     out_shape = 1
@@ -130,29 +125,13 @@ def get_1d_shape(in_shape):
     return (out_shape,)
 
 
-@constexpr
-def generate_shape_index(out_shape, indices_shape, axis):
+@_primexpr
+def generate_shape_index(out_shape, indices_shape, axis, batch_dims=0):
     out_rank = len(out_shape)
     ind_rank = len(indices_shape)
     if axis < 0:
         axis += out_rank - ind_rank + 1
-    perm_part1 = tuple(range(axis, axis + ind_rank))
+    perm_part1 = tuple(range(axis, axis + ind_rank - batch_dims))
     index = tuple(range(out_rank))
-    perm = perm_part1 + index[:axis] + index[axis + ind_rank:]
+    perm = index[:batch_dims] + perm_part1 + index[batch_dims:axis] + index[axis + ind_rank - batch_dims:]
     return perm
-
-
-@constexpr
-def is_shape_unknown(shape):
-    for i in shape:
-        if i < 0:
-            return True
-    return False
-
-
-@constexpr
-def is_dim_unknown(shape):
-    for i in shape:
-        if i == -2:
-            return True
-    return False