mindspore 2.2.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.11__cp38-cp38-manylinux1_x86_64.whl

mindspore/_extends/parallel_compile/akg_compiler/build_tbe_kernel.py

@@ -1,4 +1,4 @@
-# Copyright 2021 Huawei Technologies Co., Ltd
+# Copyright 2021-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.
@@ -25,6 +25,20 @@ from tbe.common.buildcfg import build_config
 from tbe.dsl import auto_schedule
 from tbe.dsl import build as tbe_build
 import tbe.common.context.op_context as op_context
+from impl.dynamic.add import _add_check_format, _infer_shape
+
+SHAPE = "shape"
+FORMAT = "format"
+DATA_TYPE = "data_type"
+NEW_SHAPE = "new_shape"
+ORI_SHAPE = "ori_shape"
+ORI_FORMAT = "ori_format"
+DST_TYPE = "dst_type"
+DST_ORI_SHAPE = "dst_ori_shape"
+INPUT_DESC = "input_desc"
+OUTPUT_DESC = "output_desc"
+ENABLE_VECTOR_2X = "enable_vector_2x"
+ENABLE_GROUP_INPLACE = "enable_group_inplace"
 
 
 def initialize(kernel_meta_parent_dir):
@@ -49,14 +63,14 @@ def update_config(config, op_names):
     change_type_dict = {"MatMul": (True, False),
                         "BatchMatMul": (True, False)}
     config["bool_storage_as_1bit"] = True
-    config["enable_group_inplace"] = False
-    config["enable_vector_2x"] = True
+    config[ENABLE_GROUP_INPLACE] = False
+    config[ENABLE_VECTOR_2X] = True
     for op in op_names:
         if op in bool_storage_as_1bit_oplist:
             config["bool_storage_as_1bit"] = False
         enable_group_inplace, enable_vector_2x = change_type_dict.get(op, (False, True))
-        config["enable_group_inplace"] = config["enable_group_inplace"] or enable_group_inplace
-        config["enable_vector_2x"] = config["enable_vector_2x"] and enable_vector_2x
+        config[ENABLE_GROUP_INPLACE] = config[ENABLE_GROUP_INPLACE] or enable_group_inplace
+        config[ENABLE_VECTOR_2X] = config[ENABLE_VECTOR_2X] and enable_vector_2x
 
 
 def add_new_shape(names, shapes, new_shapes, inputs):
@@ -70,11 +84,11 @@ def add_new_shape(names, shapes, new_shapes, inputs):
             continue
         if name not in inputs:
             raise RuntimeError("Can not support reshape on output tensor {}".format(name))
-        if "new_shape" not in inputs[name]:
-            inputs[name]["new_shape"] = new_shapes[i]
-        elif new_shapes[i] != inputs[name]["new_shape"]:
+        if NEW_SHAPE not in inputs[name]:
+            inputs[name][NEW_SHAPE] = new_shapes[i]
+        elif new_shapes[i] != inputs[name][NEW_SHAPE]:
             raise RuntimeError("Find different new_shape {} and {} for {}"
-                               .format(inputs[name]["new_shape"], new_shapes[i], name))
+                               .format(inputs[name][NEW_SHAPE], new_shapes[i], name))
 
 
 class TransShape:
@@ -93,22 +107,21 @@ class TransShape:
             if v.get("value") is not None:
                 continue
             names.append(k)
-            shapes.append(v["shape"])
-            ori_shapes.append(v["ori_shape"] if v.get("ori_shape") else None)
-            formats.append(v["format"])
-            ori_formats.append(v["ori_format"])
+            shapes.append(v[SHAPE])
+            ori_shapes.append(v[ORI_SHAPE] if v.get(ORI_SHAPE) else None)
+            formats.append(v[FORMAT])
+            ori_formats.append(v[ORI_FORMAT])
         if len(shapes) == 2 and len(shapes[0]) != len(shapes[1]):
-            from impl.dynamic.add import _add_check_format, _infer_shape
-            format_pattern = _add_check_format({"shape": shapes[0], "format": formats[0]},
-                                               {"shape": shapes[1], "format": formats[1]})
+            format_pattern = _add_check_format({SHAPE: shapes[0], FORMAT: formats[0]},
+                                               {SHAPE: shapes[1], FORMAT: formats[1]})
             ori_shape0 = ori_shapes[0] if ori_shapes[0] is not None else infer_ori_shape(
                 shapes[0], formats[0], ori_formats[0])
             ori_shape1 = ori_shapes[1] if ori_shapes[1] is not None else infer_ori_shape(
                 shapes[1], formats[1], ori_formats[1])
             new_shapes = [None, None]
             new_shapes[0], new_shapes[1] = _infer_shape(format_pattern,
-                                                        {"shape": shapes[0], "ori_shape": ori_shape0},
-                                                        {"shape": shapes[1], "ori_shape": ori_shape1})
+                                                        {SHAPE: shapes[0], ORI_SHAPE: ori_shape0},
+                                                        {SHAPE: shapes[1], ORI_SHAPE: ori_shape1})
             new_shapes[0], new_shapes[1], _ = shape_util.broadcast_shapes(new_shapes[0], new_shapes[1],
                                                                           param_name_input1="input0",
                                                                           param_name_input2="input1")
@@ -119,7 +132,7 @@ class TransShape:
         """deal with batch_matmul."""
         for k, v in op_inputs.items():
             # batch dimension of BatchMatMul must be fused to 1D
-            shape = v["shape"]
+            shape = v[SHAPE]
             if len(shape) > 5:
                 new_shape = [functools.reduce(lambda x, y: x * y, shape[:-4])] + shape[-4:]
                 add_new_shape(k, shape, new_shape, inputs)
@@ -135,6 +148,10 @@ class TransShape:
 
 def infer_ori_shape(shape, cur_format, ori_format):
     """Given current format and shape, infer the shape with ori_format."""
+
+    def _shape_error(current_shape, current_format):
+        raise ValueError("Invalid shape {} for format {}".format(current_shape, current_format))
+
     if cur_format == ori_format:
         return shape
     default_formats = ["DefaultFormat", "ND", "NCHW"]
@@ -145,7 +162,7 @@ def infer_ori_shape(shape, cur_format, ori_format):
     if cur_format == "FRACTAL_NZ" and ori_format in default_formats:
         dims = len(shape)
         if dims < 4:
-            raise ValueError("Invalid shape {} for format {}".format(shape, cur_format))
+            _shape_error(shape, cur_format)
         ori_shape = shape[:dims - 4]
         m = shape[-3] * shape[-2]
         n = shape[-4] * shape[-1]
@@ -155,13 +172,13 @@ def infer_ori_shape(shape, cur_format, ori_format):
 
     if cur_format == "NC1HWC0" and ori_format in default_formats:
         if len(shape) != 5:
-            raise ValueError("Invalid shape {} for format {}".format(shape, cur_format))
+            _shape_error(shape, cur_format)
         ori_shape = [shape[0], shape[1] * shape[4], shape[2], shape[3]]
         return ori_shape
 
     if cur_format == "NHWC" and ori_format in default_formats:
         if len(shape) != 4:
-            raise ValueError("Invalid shape {} for format {}".format(shape, cur_format))
+            _shape_error(shape, cur_format)
         ori_shape = [shape[0], shape[3], shape[1], shape[2]]
         return ori_shape
 
@@ -202,7 +219,7 @@ def get_input_desc(input_desc):
     res = {}
     for desc in input_desc:
         for item in desc:
-            item["shape"] = [1] if not item["shape"] else item["shape"]
+            item[SHAPE] = [1] if not item[SHAPE] else item[SHAPE]
             res[item["tensor_name"]] = item
     return res
 
@@ -215,7 +232,7 @@ def get_inputs_tensor(input_desc, all_tensors):
             name = item["tensor_name"]
             if item.get("value") is not None:
                 # const value
-                all_tensors[name] = tvm.const(item["value"], item["data_type"])
+                all_tensors[name] = tvm.const(item["value"], item[DATA_TYPE])
             if all_tensors.get(name) is None:
                 raise ValueError("Tensor [{}] not found.".format(name))
             inputs.append(all_tensors[name])
@@ -237,17 +254,17 @@ def get_op_attrs(op, fusion_op_name):
     op_name = op["name"]
     op_attrs = get_attr_dict(op.get("attr"))
     if op_name == "BatchMatMul":
-        op_attrs["dst_type"] = op["output_desc"][0]["data_type"]
-        op_attrs["dst_ori_shape"] = op["output_desc"][0].get("ori_shape")
-        if op_attrs.get("dst_ori_shape") is None:
-            op_attrs["dst_ori_shape"] = infer_ori_shape(op["output_desc"][0]["shape"],
-                                                        op["output_desc"][0]["format"],
-                                                        op["output_desc"][0]["ori_format"])
+        op_attrs[DST_TYPE] = op[OUTPUT_DESC][0][DATA_TYPE]
+        op_attrs[DST_ORI_SHAPE] = op[OUTPUT_DESC][0].get(ORI_SHAPE)
+        if op_attrs.get(DST_ORI_SHAPE) is None:
+            op_attrs[DST_ORI_SHAPE] = infer_ori_shape(op[OUTPUT_DESC][0][SHAPE],
+                                                      op[OUTPUT_DESC][0][FORMAT],
+                                                      op[OUTPUT_DESC][0][ORI_FORMAT])
     elif op_name == "MatMul":
-        op_attrs["dst_type"] = op["output_desc"][0]["data_type"]
-        op_attrs["dst_format"] = op["output_desc"][0]["format"]
+        op_attrs[DST_TYPE] = op[OUTPUT_DESC][0][DATA_TYPE]
+        op_attrs["dst_format"] = op[OUTPUT_DESC][0][FORMAT]
     elif op_name == "Cast":
-        op_attrs["dst_type"] = op["output_desc"][0]["data_type"]
+        op_attrs[DST_TYPE] = op[OUTPUT_DESC][0][DATA_TYPE]
     op_attrs["fusion_op_name"] = fusion_op_name
     return op_attrs
 
@@ -256,17 +273,17 @@ def create_placeholders(inputs):
     """Create placeholders."""
     tensors = {}
     for k, v in inputs.items():
-        dtype = v["data_type"]
+        dtype = v[DATA_TYPE]
         if dtype == "bool":
             dtype = "int8"
-        shape = v["shape"]
-        if "new_shape" in v:
-            shape = v["new_shape"]
+        shape = v[SHAPE]
+        if NEW_SHAPE in v:
+            shape = v[NEW_SHAPE]
         attr = {
-            "format": v.get("format"),
+            FORMAT: v.get(FORMAT),
             "sub_format": v.get("sub_format", ""),
-            "ori_shape": v.get("ori_shape"),
-            "ori_format": v.get("ori_format"),
+            ORI_SHAPE: v.get(ORI_SHAPE),
+            ORI_FORMAT: v.get(ORI_FORMAT),
             "addr_type": v.get("addr_type", 0),
             "valid_shape": v.get("valid_shape", []),
             "slice_offset": v.get("slice_offset", []),
@@ -276,8 +293,8 @@ def create_placeholders(inputs):
             "L1_valid_size": v.get("L1_valid_size", -1),
             "range": v.get("range", [])
         }
-        if attr.get("ori_shape") is None:
-            attr["ori_shape"] = infer_ori_shape(v.get("shape"), v.get("format"), attr.get("ori_format"))
+        if attr.get(ORI_SHAPE) is None:
+            attr[ORI_SHAPE] = infer_ori_shape(v.get(SHAPE), v.get(FORMAT), attr.get(ORI_FORMAT))
         tensors[k] = tvm.placeholder(shape=shape, name=k, dtype=dtype, attrs=attr)
     return tensors
 
@@ -289,8 +306,8 @@ def same_shape(inputs):
     base_shape = -1
     for _, v in inputs.items():
         if base_shape == -1:
-            base_shape = v["shape"]
-        if v["shape"] != base_shape:
+            base_shape = v[SHAPE]
+        if v[SHAPE] != base_shape:
             return False
     return True
 
@@ -298,17 +315,17 @@ def same_shape(inputs):
 def create_input_tensors(json_dict):
     """Create input placeholders."""
     fold_dim = True
-    inputs = get_input_desc(json_dict.get("input_desc", []))
+    inputs = get_input_desc(json_dict.get(INPUT_DESC, []))
     for op in json_dict["op_desc"]:
         op_name = op["name"]
         pattern = get_op_reg_info(op_name, "pattern")
-        op_inputs = get_input_desc(op.get("input_desc", []))
+        op_inputs = get_input_desc(op.get(INPUT_DESC, []))
         TransShape.run(op_name, pattern, op_inputs, inputs)
         if pattern != OpPattern.ELEMWISE or not same_shape(op_inputs):
             fold_dim = False
     if fold_dim:
         for k, v in inputs.items():
-            shape = v["shape"]
+            shape = v[SHAPE]
             new_shape = [functools.reduce(lambda x, y: x * y, shape[:])]
             add_new_shape(k, shape, new_shape, inputs)
     return create_placeholders(inputs)
@@ -324,28 +341,28 @@ def create_fusion_op_name(op_names):
 
 
 def update_format(json_dict):
-    """Some format like DefaultFormat is not recognized in TBE, need to covert these formats."""
+    """Some format like DefaultFormat is not recognized in TBE, need to convert these formats."""
 
     def _update_input_format(input_desc):
         for desc in input_desc:
             for item in desc:
-                if item["format"] == "DefaultFormat":
-                    item["format"] = "ND"
-                if item.get("ori_format") is None or item["ori_format"] == "DefaultFormat":
-                    item["ori_format"] = "NCHW"
+                if item[FORMAT] == "DefaultFormat":
+                    item[FORMAT] = "ND"
+                if item.get(ORI_FORMAT) is None or item[ORI_FORMAT] == "DefaultFormat":
+                    item[ORI_FORMAT] = "NCHW"
 
     def _update_output_format(output_desc):
         for item in output_desc:
-            if item["format"] == "DefaultFormat":
-                item["format"] = "ND"
-            if item.get("ori_format") is None or item["ori_format"] == "DefaultFormat":
-                item["ori_format"] = "NCHW"
+            if item[FORMAT] == "DefaultFormat":
+                item[FORMAT] = "ND"
+            if item.get(ORI_FORMAT) is None or item[ORI_FORMAT] == "DefaultFormat":
+                item[ORI_FORMAT] = "NCHW"
 
-    _update_input_format(json_dict.get("input_desc", []))
-    _update_output_format(json_dict["output_desc"])
+    _update_input_format(json_dict.get(INPUT_DESC, []))
+    _update_output_format(json_dict[OUTPUT_DESC])
     for op in json_dict["op_desc"]:
-        _update_input_format(op.get("input_desc", []))
-        _update_output_format(op["output_desc"])
+        _update_input_format(op.get(INPUT_DESC, []))
+        _update_output_format(op[OUTPUT_DESC])
 
 
 def gen_args_remap(orig_inputs_name, orig_outputs_name, inputs_name, outputs_name, inplace_names):
@@ -441,8 +458,8 @@ def build(json_str, kernel_meta_parent_dir):
     """Build kernel."""
     json_dict = json.loads(json_str)
     update_format(json_dict)
-    inputs_name = get_inputs_name(json_dict.get("input_desc", []))
-    outputs_name, inplace_names = get_outputs_info(json_dict["output_desc"])
+    inputs_name = get_inputs_name(json_dict.get(INPUT_DESC, []))
+    outputs_name, inplace_names = get_outputs_info(json_dict[OUTPUT_DESC])
     op_names = get_all_op_name(json_dict["op_desc"])
     fusion_op_name = create_fusion_op_name(op_names)
 
@@ -458,7 +475,7 @@ def build(json_str, kernel_meta_parent_dir):
         for op in json_dict["op_desc"]:
             op_name = op["name"]
             # get op input tensor
-            op_inputs = get_inputs_tensor(op.get("input_desc", []), all_tensors)
+            op_inputs = get_inputs_tensor(op.get(INPUT_DESC, []), all_tensors)
             # get op attrs
             op_attrs = get_op_attrs(op, fusion_op_name)
             # op compute
@@ -466,10 +483,10 @@ def build(json_str, kernel_meta_parent_dir):
             # update op output tensor
             if not isinstance(op_outputs, (list, tuple)):
                 op_outputs = [op_outputs]
-            if len(op["output_desc"]) != len(op_outputs):
+            if len(op[OUTPUT_DESC]) != len(op_outputs):
                 raise ValueError("len(op[\"output_desc\"] is not equal to the number of real output tensors in op[{}]: "
-                                 "{} vs {}".format(op_name, len(op["output_desc"]), len(op_outputs)))
-            for i, desc in enumerate(op["output_desc"]):
+                                 "{} vs {}".format(op_name, len(op[OUTPUT_DESC]), len(op_outputs)))
+            for i, desc in enumerate(op[OUTPUT_DESC]):
                 all_tensors[desc["tensor_name"]] = op_outputs[i]
 
         # Collect input, output tensors

mindspore/_extends/parallel_compile/akg_compiler/tbe_topi.py

@@ -310,8 +310,8 @@ def _log(x, attrs=None):
     if base <= 0 and not math.isclose(base, -1.0, rel_tol=1e-8, abs_tol=0.0):
         raise ValueError("base must be strictly positive or -1, but got {}".format(base))
     from impl.log import log_compute
-    outputDesc = {"dtype": x.dtype, "shape": x.shape}
-    return log_compute(x, outputDesc, base, scale, shift, kernel_name=attrs["fusion_op_name"])
+    output_desc = {"dtype": x.dtype, "shape": x.shape}
+    return log_compute(x, output_desc, base, scale, shift, kernel_name=attrs["fusion_op_name"])
 
 
 @reg_op("Maximum", pattern=OpPattern.ELEMWISE)
@@ -350,8 +350,8 @@ def _mul(x0, x1, attrs=None):
         return tbe.dsl.vmuls(x1, x0)
     x0, x1 = _broadcast(x0, x1)
     from impl.mul import mul_compute
-    outputDesc = {"dtype": x0.dtype, "shape": x0.shape}
-    return mul_compute(x0, x1, outputDesc, kernel_name=attrs["fusion_op_name"])
+    output_desc = {"dtype": x0.dtype, "shape": x0.shape}
+    return mul_compute(x0, x1, output_desc, kernel_name=attrs["fusion_op_name"])
 
 
 @reg_op("Neg", pattern=OpPattern.ELEMWISE)

mindspore/_extends/parallel_compile/akg_compiler/util.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Huawei Technologies Co., Ltd
+# Copyright 2022-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.
@@ -19,6 +19,9 @@ import shutil
 from mindspore import log as logger
 from mindspore._extends.parallel_compile.akg_compiler.tbe_topi import get_op_reg_info
 
+O_SUFFIX = ".o"
+JSON_SUFFIX = ".json"
+
 
 def update_attr(attr, new_attr):
     """Update new_attr to attr."""
@@ -111,8 +114,8 @@ def select_best(src_dirs, dst_dir, op_name):
     max_block_dim = 1
     max_block_dim_idx = -1
     for i, src_dir in enumerate(src_dirs):
-        o_path = os.path.join(src_dir, op_name + ".o")
-        json_path = os.path.join(src_dir, op_name + ".json")
+        o_path = os.path.join(src_dir, op_name + O_SUFFIX)
+        json_path = os.path.join(src_dir, op_name + JSON_SUFFIX)
         if os.path.isfile(o_path) and os.path.isfile(json_path):
             with open(json_path, 'r') as f:
                 json_str = f.read()
@@ -121,10 +124,10 @@ def select_best(src_dirs, dst_dir, op_name):
                     max_block_dim_idx = i
                     max_block_dim = json_dict["blockDim"]
     if max_block_dim_idx >= 0:
-        o_path = os.path.join(src_dirs[max_block_dim_idx], op_name + ".o")
-        json_path = os.path.join(src_dirs[max_block_dim_idx], op_name + ".json")
-        _copy_file(o_path, os.path.join(dst_dir, op_name + ".o"))
-        _copy_file(json_path, os.path.join(dst_dir, op_name + ".json"))
+        o_path = os.path.join(src_dirs[max_block_dim_idx], op_name + O_SUFFIX)
+        json_path = os.path.join(src_dirs[max_block_dim_idx], op_name + JSON_SUFFIX)
+        _copy_file(o_path, os.path.join(dst_dir, op_name + O_SUFFIX))
+        _copy_file(json_path, os.path.join(dst_dir, op_name + JSON_SUFFIX))
         logger.info("{}, best compile result dir: {}".format(op_name, src_dirs[max_block_dim_idx]))
         return True
     logger.info("{}, best compile result dir not found".format(op_name))

mindspore/_extends/parallel_compile/tbe_compiler/tbe_helper.py

@@ -65,7 +65,8 @@ class LocalLock:
 
     def __init__(self, lock_file):
         if not os.path.exists(lock_file):
-            write_to_file(lock_file)
+            if not write_to_file(lock_file):
+                raise IOError("write_to_file failed.")
         self.lock_fd = os.open(lock_file, os.O_WRONLY | os.O_CREAT, stat.S_IWUSR | stat.S_IRUSR)
 
     def __del__(self):

mindspore/_extends/parse/__init__.py

@@ -28,7 +28,8 @@ from .parser import (Parser, create_instance, is_supported_create_instance_type,
                      convert_class_to_function, convert_cell_list_to_sequence, is_cell_list, get_obj_from_sequence,
                      get_type, is_class_member_recursive, get_global_params, get_adapter_tensor_attr,
                      get_obj_defined_from_obj_type, is_from_third_party_library, get_const_abs, get_const_round,
-                     get_const_len, is_adapter_tensor_class, is_adapter_parameter_class, convert_to_namedtuple)
+                     get_const_len, is_adapter_tensor_class, is_adapter_parameter_class, convert_to_namedtuple,
+                     is_module_list)
 
 __all__ = ['Parser', 'create_instance', 'is_supported_create_instance_type', 'generate_scope', 'get_attr_from_object',
            'get_bprop_method_of_class', 'get_class_instance_type', 'get_class_member_namespace_symbol',
@@ -41,4 +42,4 @@ __all__ = ['Parser', 'create_instance', 'is_supported_create_instance_type', 'ge
            'convert_class_to_function', 'convert_cell_list_to_sequence', 'is_cell_list', 'get_obj_from_sequence',
            'get_type', 'is_class_member_recursive', 'get_adapter_tensor_attr', 'get_obj_defined_from_obj_type',
            'is_from_third_party_library', 'get_const_abs', 'get_const_round', 'get_const_len',
-           'is_adapter_tensor_class', 'is_adapter_parameter_class', 'convert_to_namedtuple',]
+           'is_adapter_tensor_class', 'is_adapter_parameter_class', 'convert_to_namedtuple', 'is_module_list']

mindspore/_extends/parse/parser.py

@@ -493,9 +493,14 @@ def is_cell_list(obj):
     return isinstance(obj, nn.CellList)
 
 
+def is_module_list(obj):
+    """Check if obj is nn.ModuleList"""
+    return hasattr(obj, "__cell_as_list__") and not isinstance(obj, nn.CellList)
+
+
 def convert_cell_list_to_sequence(obj):
     """Convert nn.CellList to sequence."""
-    if not isinstance(obj, nn.CellList):
+    if not hasattr(obj, "__cell_as_list__"):
         raise TypeError(f"Obj should be nn.CellList, but got {obj}")
     if not hasattr(obj, "_cells"):
         raise AttributeError(f"nn.CellList is missing _cells property.")

mindspore/_extends/parse/standard_method.py

@@ -2349,12 +2349,22 @@ def ms_iter(xs):
 
 def ms_next(it):
     """Implementation of `next`."""
-    return it.__ms_next__()
+    return it.__ms_next__
 
 
 def hasnext(it):
     """Implementation of `hasnext`."""
-    return it.__ms_hasnext__()
+    return it.__ms_hasnext__
+
+
+def str_next(xs):
+    """Next string."""
+    return xs[0], xs[1:]
+
+
+def str_hasnext(xs):
+    """Whether the string is empty or not."""
+    return len(xs) > 0
 
 
 @constexpr
@@ -2819,6 +2829,8 @@ def enumerate_(x, start=0):
     x_type = F.typeof(x)
     ret = ()
     op_name = "enumerate"
+    if isinstance(x, (int, float, bool)):
+        raise TypeError(f"For 'enumerate', the 'first input' should be tuple or list or tensor, but got {type(x)}.")
     if check_is_const_int(start, op_name, "start"):
         if check_is_tensor(x_type):
             for i in range(x.shape[0]):
@@ -3254,15 +3266,6 @@ def check_is_tensor(x):
     return False
 
 
-@constexpr
-def check_is_tuple_or_list_or_tensor(x, op_name, arg_name):
-    """check whether x is list or tuple or tensor."""
-    if isinstance(x, (mstype.List, mstype.Tuple, mstype.TensorType)):
-        return True
-    raise TypeError(
-        f"For '{op_name}', the '{arg_name}' should be tuple or list or tensor, but got {x}.")
-
-
 def check_is_const_int(x, op_name, arg_name):
     """check whether x is const int."""
     if x is None:

mindspore/_extends/remote/kernel_build_server.py

@@ -15,7 +15,8 @@
 """kernel build server"""
 import os
 from mindspore import log as logger
-from mindspore._extends.parallel_compile.akg_compiler.akg_process import create_akg_parallel_process, create_akg_v2_parallel_process
+from mindspore._extends.parallel_compile.akg_compiler.akg_process import create_akg_parallel_process, \
+    create_akg_v2_parallel_process
 
 
 class Messager:

mindspore/common/_utils.py

@@ -116,3 +116,19 @@ def load_lib(lib_path):
         logger.warning(f'Loading {lib_path} lib error.')
         return False
     return True
+
+
+def _jit_fallback_next_func(xs):
+    """Generate ms_next for xs"""
+    if hasattr(xs, "__next__"):
+        # Convert an iterator to tuple first.
+        xs = tuple(xs)
+    return xs[0], xs[1:]
+
+
+def _jit_fallback_has_next_func(xs):
+    """Determine whether xs has next value"""
+    if hasattr(xs, "__next__"):
+        # Convert an iterator to tuple first.
+        xs = tuple(xs)
+    return len(xs) > 0

mindspore/common/api.py

@@ -1060,7 +1060,7 @@ def _build_broadcast_graph(broadcast_params_dict, broadcast_phase):
     _broadcast_net.phase = broadcast_phase
     broadcasted_params = _broadcast_net()
     for param_name, param in zip(broadcast_params_dict.keys(), broadcasted_params):
-        broadcast_params_dict[param_name].set_data(param)
+        broadcast_params_dict.get(param_name).set_data(param)
 
 
 def _get_auto_split_param_names(parameter_layout_dict):