mindspore 2.1.0__cp37-cp37m-manylinux1_x86_64.whl → 2.2.11__cp37-cp37m-manylinux1_x86_64.whl

mindspore/ops/operations/custom_ops.py

@@ -42,6 +42,24 @@ from ._pyfunc_registry import add_pyfunc
 if platform.system() != "Windows":
     import fcntl
 
+KEY_ATTR = "attr"
+KEY_NAME = "name"
+INPUT_NAMES = "input_names"
+ATTR_NAMES = "attr_names"
+AUTO_DIFF = "autodiff"
+IMPLY_TYPE = "imply_type"
+FUSION_TYPE = "fusion_type"
+MS_KERNEL_FLAG = "ms_kernel_flag"
+AKG = "AKG"
+TBE = "TBE"
+CUDA = "CUDA"
+AICORE = "AiCore"
+CPU = "CPU"
+GPU = "GPU"
+ASCEND = "Ascend"
+HYBRID_TYPE = "hybrid"
+OP_NAME = "op_name"
+
 
 def _get_cache_path():
     """
@@ -146,10 +164,10 @@ class Custom(ops.PrimitiveWithInfer):
     function if needed. Then these `Custom` objects can be directly used in neural networks.
     Detailed description and introduction of user-defined operators, including correct writing of parameters,
     please refer to `Custom Operators Tutorial
-    <https://www.mindspore.cn/tutorials/experts/en/r2.1/operation/op_custom.html>`_ .
+    <https://www.mindspore.cn/tutorials/experts/en/r2.2/operation/op_custom.html>`_ .
 
     .. warning::
-        This is an experimental API that is subject to change.
+        - This is an experimental API that is subject to change.
 
     .. note::
         The supported platforms are determined by the input `func_type`. The supported platforms are as follows:
@@ -162,6 +180,12 @@ class Custom(ops.PrimitiveWithInfer):
         - "julia": supports ["CPU"].
         - "aicpu": supports ["Ascend"].
 
+        If run on ge backend, use `CustomRegOp` to generate the registration information of "aicpu" and "tbe" operator,
+        use `custom_info_register` to bind the registration information to the `func` of the "tbe" operator,
+        then save the registration information of "aicpu" operator and the `func` implementation of "tbe" operator to
+        a file or separate files, keep these files in a separate directory, and set the absolute path of this directory
+        to environment variable "MS_DEV_CUSTOM_OPP_PATH" before running the network.
+
     Args:
         func (Union[function, str]):
 
@@ -446,10 +470,10 @@ class Custom(ops.PrimitiveWithInfer):
     op_path_in_cache = []  # Save paths for op functions created in the cached.
     custom_aot_warning = True  # Flag to enable warnings about custom aot path white list
 
-    def __init__(self, func, out_shape=None, out_dtype=None, func_type="hybrid", bprop=None, reg_info=None):
-        ops.PrimitiveWithInfer.__init__(self, "Custom")
+    def __init__(self, func, out_shape=None, out_dtype=None, func_type=HYBRID_TYPE, bprop=None, reg_info=None):
+        super().__init__("Custom")
 
-        self.supported_targets = ["Ascend", "GPU", "CPU"]
+        self.supported_targets = [ASCEND, GPU, CPU]
         self.supported_func_type = ["hybrid", "akg", "tbe", "aicpu", "aot", "pyfunc", "julia"]
         self.log_prefix = "For '{}', 'func_type': {}, 'func': {}".format(self.name, func_type, func)
         self.func = func
@@ -466,7 +490,7 @@ class Custom(ops.PrimitiveWithInfer):
         self._update_func_info(reg_info)
         self.add_prim_attr("func_name", self.func_name)
         self.add_prim_attr("uniq_name", self.uniq_name)
-        if self.func_type == "hybrid":
+        if self.func_type == HYBRID_TYPE:
             self.add_prim_attr("func_compile_attrs", self._func_compile_attrs)
 
         self.add_prim_attr("imply_path", self.imply_path)
@@ -495,7 +519,7 @@ class Custom(ops.PrimitiveWithInfer):
         if func_type == "akg":
             self._set_akg_kernel_type()
 
-        if not self.bprop and self.func_type == "hybrid":
+        if not self.bprop and self.func_type == HYBRID_TYPE:
             self._hybrid_autodiff(func_type)
 
         self.add_prim_attr("func_type", self.func_type)
@@ -570,7 +594,7 @@ class Custom(ops.PrimitiveWithInfer):
         elif "compute" in self.func_source_str:
             self.func_type = "tvm_compute"
         else:
-            self.func_type = "hybrid"
+            self.func_type = HYBRID_TYPE
             self._hybrid_func_analyser()
 
     def _check_julia_func(self):
@@ -620,24 +644,24 @@ class Custom(ops.PrimitiveWithInfer):
                     raise TypeError(
                         "{}, the legal path for the file is {}, but the file is {}".format(
                             self.log_prefix, legal_path, file_path))
-            if not file_path.endswith("so"):
+            if file_path.endswith(("cu", "cpp", "cc")):
                 file_path = _compile_aot(file_path)
                 self.func = file_path + ":" + file_name_list[1]
 
         elif self.func_type == "julia":
             self._check_julia_func()
-        elif self.func_type == "hybrid":
-            if not hasattr(self.func, "ms_kernel_flag"):
+        elif self.func_type == HYBRID_TYPE:
+            if not hasattr(self.func, MS_KERNEL_FLAG):
                 raise TypeError("{}, 'func' must be a function decorated by kernel".format(self.log_prefix))
             self._is_ms_kernel = True
             self._func_compile_attrs = getattr(self.func, "compile_attrs", {})
         elif self.func_type == "akg":
-            if hasattr(self.func, "ms_kernel_flag"):
+            if hasattr(self.func, MS_KERNEL_FLAG):
                 logger.warning("{}. To have a better user experience, the mode hybrid is suggested "
                                "for the input function with decorator @kernel. "
                                "To enable this mode, set the 'func_type' to be \"hybrid\"".format(self.log_prefix))
         elif self.func_type == "pyfunc":
-            if hasattr(self.func, "ms_kernel_flag"):
+            if hasattr(self.func, MS_KERNEL_FLAG):
                 logger.warning("{}. Now you are using the function with decorator @kernel in the mode pyfunc. "
                                "The kernel will be executed as a native python function, which might lead to "
                                "low efficiency. To accelerate the kernel, set the 'func_type' to be \"hybrid\""
@@ -751,7 +775,7 @@ class Custom(ops.PrimitiveWithInfer):
                     continue
                 if isinstance(reg_info_item, str):
                     reg_info_item = json.loads(reg_info_item)
-                prefix = "_".join([prefix, reg_info_item.get("op_name", "")])
+                prefix = "_".join([prefix, reg_info_item.get(OP_NAME, "")])
             self.uniq_name = prefix + "_" + self.func_name
         else:
             raise TypeError("For '{}', 'func' must be of type function or str, but got {}"
@@ -761,23 +785,23 @@ class Custom(ops.PrimitiveWithInfer):
         """Update op attrs in reg_info."""
         output_name_list = []
         for _, item in enumerate(reg_info.get("outputs", [])):
-            if isinstance(item, dict) and item.get("name"):
-                output_name_list.append(item.get("name"))
+            if isinstance(item, dict) and item.get(KEY_NAME):
+                output_name_list.append(item.get(KEY_NAME))
         if output_name_list:
             self.add_prim_attr("output_names", output_name_list)
 
-        if isinstance(reg_info.get("op_name"), str):
-            self.add_prim_attr("reg_op_name", reg_info.get("op_name"))
+        if isinstance(reg_info.get(OP_NAME), str):
+            self.add_prim_attr("reg_op_name", reg_info.get(OP_NAME))
 
         if self.func_type == "aicpu":
-            self.uniq_name = reg_info["op_name"]
+            self.uniq_name = reg_info[OP_NAME]
             self.add_prim_attr("uniq_name", self.uniq_name)
 
         if self.func_type in ["aot", "aicpu"]:
-            if reg_info.get("attr") is not None and isinstance(reg_info["attr"], list):
-                for item in reg_info["attr"]:
+            if reg_info.get(KEY_ATTR) is not None and isinstance(reg_info[KEY_ATTR], list):
+                for item in reg_info[KEY_ATTR]:
                     if isinstance(item, dict) and item.get("value") is not None:
-                        self.add_prim_attr(item["name"], item["value"])
+                        self.add_prim_attr(item[KEY_NAME], item["value"])
 
     def _register_info(self, info):
         """Register reg_info."""
@@ -795,7 +819,7 @@ class Custom(ops.PrimitiveWithInfer):
             if isinstance(reg_info, str):
                 reg_info = json.loads(reg_info)
             if self.fake_output:
-                reg_info["outputs"].append(dict({"index": 0, "name": "y", "param_type": "required"}))
+                reg_info["outputs"].append(dict({"index": 0, KEY_NAME: "y", "param_type": "required"}))
                 new_dtype_format = []
                 for i in reg_info["dtype_format"]:
                     new_dtype_format.append(i + (DataType.I32_Default,))
@@ -867,16 +891,16 @@ class Custom(ops.PrimitiveWithInfer):
                             "'CustomRegOp' to generate the registration information, then pass it to 'reg_info' or "
                             "use 'custom_info_register' to bind it to 'func' if 'func' is a function."
                             .format(self.log_prefix, reg_info, type(reg_info)))
-        reg_info["op_name"] = self.uniq_name
-        reg_info["imply_type"] = self._get_imply_type(reg_info, target)
-        if not isinstance(reg_info.get("fusion_type"), str) or not reg_info["fusion_type"].strip():
-            reg_info["fusion_type"] = "OPAQUE"
+        reg_info[OP_NAME] = self.uniq_name
+        reg_info[IMPLY_TYPE] = self._get_imply_type(reg_info, target)
+        if not isinstance(reg_info.get(FUSION_TYPE), str) or not reg_info[FUSION_TYPE].strip():
+            reg_info[FUSION_TYPE] = "OPAQUE"
         # Supplement necessary info for TBE if these information is missing in reg_info
-        if reg_info["imply_type"] == "TBE":
-            if reg_info.get("attr") is not None and isinstance(reg_info["attr"], list):
-                for i, item in enumerate(reg_info["attr"]):
+        if reg_info[IMPLY_TYPE] == TBE:
+            if reg_info.get(KEY_ATTR) is not None and isinstance(reg_info[KEY_ATTR], list):
+                for i, item in enumerate(reg_info[KEY_ATTR]):
                     if isinstance(item, dict) and item.get("value") is None:
-                        reg_info["attr"][i]["value"] = "all"
+                        reg_info[KEY_ATTR][i]["value"] = "all"
             reg_info["async_flag"] = reg_info.get("async_flag", False)
             reg_info["binfile"] = "%s.so" % self.func_name
             reg_info["compute_cost"] = reg_info.get("compute_cost", 10)
@@ -884,8 +908,8 @@ class Custom(ops.PrimitiveWithInfer):
             reg_info["partial_flag"] = reg_info.get("partial_flag", True)
             reg_info["needCheckSupport"] = reg_info.get("need_check_supported", False)
         # Supplement necessary info for AKG if these information is missing in reg_info
-        if reg_info["imply_type"] == "AKG":
-            target_to_processor = {"Ascend": "AiCore", "GPU": "CUDA", "CPU": "CPU"}
+        if reg_info[IMPLY_TYPE] == AKG:
+            target_to_processor = {ASCEND: AICORE, GPU: CUDA, CPU: CPU}
             reg_info["processor"] = reg_info.get("processor", target_to_processor.get(target))
         return reg_info
 
@@ -898,15 +922,15 @@ class Custom(ops.PrimitiveWithInfer):
             # Infer target from reg_info["processor"], reg_info generated from AkgGpuRegOp or AkgAscendRegOp
             #   will have the processor information.
             if target not in self.supported_targets:
-                processor_to_target = {"AiCore": "Ascend", "CUDA": "GPU", "CPU": "CPU"}
+                processor_to_target = {AICORE: ASCEND, CUDA: GPU, CPU: CPU}
                 target = processor_to_target.get(reg_info.get("processor"))
-            # Infer target from reg_info["imply_type"]
+            # Infer target from reg_info[IMPLY_TYPE]
             if target not in self.supported_targets:
-                imply_type_to_target = {"TBE": "Ascend", "GPU": "GPU", "CPU": "CPU"}
-                target = imply_type_to_target.get(reg_info.get("imply_type"))
+                imply_type_to_target = {TBE: ASCEND, GPU: GPU, CPU: CPU}
+                target = imply_type_to_target.get(reg_info.get(IMPLY_TYPE))
         # Infer target from func_type
         if target not in self.supported_targets:
-            func_type_to_target = {"tbe": "Ascend", "pyfunc": "CPU"}
+            func_type_to_target = {"tbe": ASCEND, "pyfunc": CPU}
             target = func_type_to_target.get(self.func_type)
         if target not in self.supported_targets:
             raise ValueError("{}, target set in registration information must be one of {}, but got {}"
@@ -915,14 +939,14 @@ class Custom(ops.PrimitiveWithInfer):
 
     def _get_imply_type(self, reg_info, target):
         """Get imply_typ information."""
-        # Get imply_type from reg_info["imply_type"]
-        if isinstance(reg_info, dict) and isinstance(reg_info.get("imply_type"), str) and \
-                reg_info["imply_type"].strip():
-            return reg_info["imply_type"]
+        # Get imply_type from reg_info[IMPLY_TYPE]
+        if isinstance(reg_info, dict) and isinstance(reg_info.get(IMPLY_TYPE), str) and \
+                reg_info[IMPLY_TYPE].strip():
+            return reg_info[IMPLY_TYPE]
         # Infer imply_type from func_type
-        func_type_to_imply_type = {"hybrid": "AKG", "akg": "AKG", "tbe": "TBE", "aicpu": "AiCPU", "pyfunc": target,
-                                   "julia": target, "aot": "BiSheng" if target == "Ascend" else target}
-        return func_type_to_imply_type.get(self.func_type, "AKG")
+        func_type_to_imply_type = {"hybrid": AKG, "akg": AKG, "tbe": TBE, "aicpu": "AiCPU", "pyfunc": target,
+                                   "julia": target, "aot": "BiSheng" if target == ASCEND else target}
+        return func_type_to_imply_type.get(self.func_type, AKG)
 
     def _save_attr(self, reg_info):
         """Save input_names and attr_names of current func."""
@@ -936,18 +960,18 @@ class Custom(ops.PrimitiveWithInfer):
             return value
 
         tensor_inputs = _get_value_list("inputs")
-        attr = _get_value_list("attr")
+        attr = _get_value_list(KEY_ATTR)
         input_names = []  # include tensor input names and attr input names
         attr_names = []
         pure_input_names = []
         for item in tensor_inputs:
-            if isinstance(item, dict) and item.get("name") is not None:
-                input_names.append(item["name"])
-                pure_input_names.append(item["name"])
+            if isinstance(item, dict) and item.get(KEY_NAME) is not None:
+                input_names.append(item[KEY_NAME])
+                pure_input_names.append(item[KEY_NAME])
         # attr is converted from inputs only when graph mode or when inputs name is also in reg info
         attr_to_input_safe = bool(input_names) or context.get_context("mode") == ms.GRAPH_MODE
         for item in attr:
-            if isinstance(item, dict) and item.get("name") is not None:
+            if isinstance(item, dict) and item.get(KEY_NAME) is not None:
                 # for custom op with function tbe, we always add attrs to inputs as we don't
                 # deal with attr value here and leave them to the backend process to fit the
                 # usual process of tbe op compiling in mindspore
@@ -956,9 +980,9 @@ class Custom(ops.PrimitiveWithInfer):
                 # add attr name to input name only when the value of attr is None in reg info
                 # as we need to get values of attrs from inputs
                 if attr_to_input_safe and (self.func_type == "tbe" or item.get("value", None) is None):
-                    input_names.append(item["name"])
-                attr_names.append(item["name"])
-        cur_attr = {"input_names": input_names, "attr_names": attr_names, "pure_input_names": pure_input_names}
+                    input_names.append(item[KEY_NAME])
+                attr_names.append(item[KEY_NAME])
+        cur_attr = {INPUT_NAMES: input_names, ATTR_NAMES: attr_names, "pure_input_names": pure_input_names}
         # If func does not have attr, save current attr.
         # Else, check if current attr is same as previous saved one.
         prev_attr_names = attr_names
@@ -967,13 +991,13 @@ class Custom(ops.PrimitiveWithInfer):
             if not isinstance(func_attr, dict):
                 setattr(self.func, "func_attr", cur_attr)
             else:
-                prev_attr_names = func_attr.get("attr_names")
+                prev_attr_names = func_attr.get(ATTR_NAMES)
         elif isinstance(self.func, str):
             func_attr = Custom.attr_dict.get(self.func)
             if not isinstance(func_attr, dict):
                 Custom.attr_dict[self.func] = cur_attr
             else:
-                prev_attr_names = func_attr.get("attr_names")
+                prev_attr_names = func_attr.get(ATTR_NAMES)
         if attr_names != prev_attr_names:
             raise ValueError("{}, attr names set in registration information must be the same as previous saved one, "
                              "but got {} vs {}".format(self.log_prefix, attr_names, prev_attr_names))
@@ -982,23 +1006,23 @@ class Custom(ops.PrimitiveWithInfer):
         """Add primitive_target to primitive's attr."""
         registered_targets = self._get_registered_targets()
         if self.func_type == "pyfunc":
-            self.set_device("CPU")
-            if registered_targets and registered_targets != ["CPU"]:
+            self.set_device(CPU)
+            if registered_targets and registered_targets != [CPU]:
                 logger.warning("{}, only supports CPU platform, but got registered target {}. "
                                "We will run it on CPU".format(self.log_prefix, registered_targets))
         elif self.func_type == "aot":
             if len(registered_targets) != 1:
                 logger.info("{}, target will be set according to context.".format(self.log_prefix))
-            elif registered_targets == ["GPU"]:
-                self.set_device("GPU")
-            elif registered_targets == ["CPU"]:
-                self.set_device("CPU")
+            elif registered_targets == [GPU]:
+                self.set_device(GPU)
+            elif registered_targets == [CPU]:
+                self.set_device(CPU)
         elif self.func_type == "julia":
-            self.set_device("CPU")
+            self.set_device(CPU)
             device_target = context.get_context('device_target')
-            if device_target == "CPU":
+            if device_target == CPU:
                 pass
-            elif device_target == "GPU" and registered_targets and registered_targets == ["CPU"]:
+            elif device_target == GPU and registered_targets and registered_targets == [CPU]:
                 logger.warning("{}, only supports CPU platform, but got registered target {}. "
                                "We will run it on CPU".format(self.log_prefix, registered_targets))
             else:
@@ -1021,15 +1045,15 @@ class Custom(ops.PrimitiveWithInfer):
         elif isinstance(self.func, str):
             func_attr = Custom.attr_dict.get(self.func)
         if isinstance(func_attr, dict):
-            _add_prim_attr("input_names")
-            _add_prim_attr("attr_names")
+            _add_prim_attr(INPUT_NAMES)
+            _add_prim_attr(ATTR_NAMES)
             _add_prim_attr("pure_input_names")
         self._add_prim_target()
         if callable(self.func) and callable(self.out_shape):
-            if hasattr(self.out_shape, "type") and getattr(self.out_shape, "type") == "autodiff":
-                self.add_prim_attr("autodiff", True)
+            if hasattr(self.out_shape, "type") and getattr(self.out_shape, "type") == AUTO_DIFF:
+                self.add_prim_attr(AUTO_DIFF, True)
         else:
-            self.add_prim_attr("autodiff", False)
+            self.add_prim_attr(AUTO_DIFF, False)
 
     def _hybrid_autodiff(self, input_func_type):
         """generate backward op for a custom hybrid op"""
@@ -1045,7 +1069,7 @@ class Custom(ops.PrimitiveWithInfer):
             def infer_func(*args):
                 return args[:inputs_num]
 
-            setattr(infer_func, "type", "autodiff")
+            setattr(infer_func, "type", AUTO_DIFF)
             op = Custom(func=self.func, out_shape=infer_func, out_dtype=infer_func,
                         func_type=input_func_type, bprop=True)
             self.bprop = grad_func(op)

mindspore/ops/operations/debug_ops.py

@@ -58,7 +58,7 @@ class ScalarSummary(Primitive):
     This operator will put a scalar to a summary file with protocol buffer format. It must be used with SummaryRecord
     or SummaryCollector, which specify the directory of the summary file. The summary file can
     be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
-    mindinsight/docs/en/r2.1/index.html>`_ for details.
+    mindinsight/docs/en/r2.2/index.html>`_ for details.
 
     Inputs:
         - **name** (str) - The name of the input variable, it must not be an empty string.
@@ -104,6 +104,7 @@ class ScalarSummary(Primitive):
             raise ValueError('The Summary is not supported, please without `-s on` and recompile source.')
 
         self.add_prim_attr("side_effect_io", True)
+        self.add_prim_attr("channel_name", "ms_scalar_summary")
 
     def __call__(self, *args):
         _cache_summary_data(self.name, args[0], args[1])
@@ -114,7 +115,7 @@ class ImageSummary(PrimitiveWithInfer):
     This operator will put an image tensor to a summary file with protocol buffer format. It must be used with
     SummaryRecord or SummaryCollector, which specify the directory of the summary file. The summary file can
     be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
-    mindinsight/docs/en/r2.1/index.html>`_ for details.
+    mindinsight/docs/en/r2.2/index.html>`_ for details.
 
     Inputs:
         - **name** (str) - The name of the input variable, it must not be an empty string.
@@ -153,6 +154,7 @@ class ImageSummary(PrimitiveWithInfer):
             raise ValueError('The Summary is not supported, please without `-s on` and recompile source.')
 
         self.add_prim_attr("side_effect_io", True)
+        self.add_prim_attr("channel_name", "ms_image_summary")
 
     def __infer__(self, name, value):
         _check_summary_param(name, value, self.__class__.__name__)
@@ -175,7 +177,7 @@ class TensorSummary(Primitive):
     This operator will put a tensor to a summary file with protocol buffer format. It must be used with SummaryRecord
     or SummaryCollector, which specify the directory of the summary file. The summary file can
     be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
-    mindinsight/docs/en/r2.1/index.html>`_ for details.
+    mindinsight/docs/en/r2.2/index.html>`_ for details.
 
     Inputs:
         - **name** (str) - The name of the input variable.
@@ -221,6 +223,7 @@ class TensorSummary(Primitive):
             raise ValueError('The Summary is not supported, please without `-s on` and recompile source.')
 
         self.add_prim_attr("side_effect_io", True)
+        self.add_prim_attr("channel_name", "ms_tensor_summary")
 
     def __call__(self, *args):
         _cache_summary_data(self.name, args[0], args[1])
@@ -231,7 +234,7 @@ class HistogramSummary(PrimitiveWithInfer):
     This operator will calculate the histogram of a tensor and put it to a summary file with protocol buffer format.
     It must be used with SummaryRecord or SummaryCollector, which specify the directory of the summary file.
     The summary file can be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
-    mindinsight/docs/en/r2.1/index.html>`_ for details.
+    mindinsight/docs/en/r2.2/index.html>`_ for details.
 
     Inputs:
         - **name** (str) - The name of the input variable.
@@ -276,6 +279,7 @@ class HistogramSummary(PrimitiveWithInfer):
             raise ValueError('The Summary is not supported, please without `-s on` and recompile source.')
 
         self.add_prim_attr("side_effect_io", True)
+        self.add_prim_attr("channel_name", "ms_histogram_summary")
 
     def __infer__(self, name, value):
         _check_summary_param(name, value, self.__class__.__name__)

mindspore/ops/operations/image_ops.py

@@ -180,7 +180,7 @@ class AdjustHue(Primitive):
 
 
 class ExtractGlimpse(Primitive):
-    """
+    r"""
     Extracts glimpses(usually subarea of rectangle) from the input image Tensor and return as windows.
 
     Note:
@@ -205,20 +205,20 @@ class ExtractGlimpse(Primitive):
             - When `noise` is ``'zero'`` , the value of `uniform_noise` must be ``'False'`` and the
               filling noise will be zero so that the result is fixed.
             - When `uniform_noise` is ``'True'`` , the value of `noise` only can be ``'uniform'`` .
-              When `uniform_noise` is ``'False'`` , the value of `noise` can be ``'uniform'`` , ``'gaussian'`` and
+              When `uniform_noise` is ``'False'`` , the value of `noise` can be ``'uniform'`` , ``'gaussian'`` or
               ``'zero'`` .
 
     Inputs:
-        - **x** (Tensor) - A 4-D float tensor of shape :math:`(batch_size, height, width, channels)`.
+        - **x** (Tensor) - A 4-D float tensor of shape :math:`(batch\_size, height, width, channels)`.
           Types allowed: float32.
         - **size** (Tensor) - A 1-D tensor of 2 elements containing the size of the glimpses to extract.
           The glimpse height must be specified first, following by the glimpse width. Types allowed: int32.
           The value of size must be greater than zero.
-        - **offsets** (Tensor) - A 2-D integer tensor of shape :math:`(batch_size, 2)` containing the y, x locations
+        - **offsets** (Tensor) - A 2-D integer tensor of shape :math:`(batch\_size, 2)` containing the y, x locations
           of the center of each window. Types allowed: float32.
 
     Outputs:
-        A 4-D tensor of shape :math:`(batch_size, glimpse_height, glimpse_width, channels)` with type: float32.
+        A 4-D tensor of shape :math:`(batch\_size, glimpse\_height, glimpse\_width, channels)` with type: float32.
 
     Raises:
         TypeError: If `centered` is not a bool.
@@ -277,9 +277,16 @@ class CropAndResize(Primitive):
 
     Args:
         method (str, optional): An optional string that specifies the sampling method for resizing.
-            It can be ``"bilinear"`` , ``"nearest"`` or ``"bilinear_v2"`` . The option "bilinear" stands for standard
-            ``"bilinear"`` interpolation algorithm, while ``"bilinear_v2"`` may result in better result in some cases.
-            Default: ``"bilinear"`` .
+            It can be ``"bilinear"`` , ``"nearest"`` or ``"bilinear_v2"`` . Default: ``"bilinear"`` .
+
+            - ``"nearest"``: Nearest neighbor interpolation. Each output pixel is assigned the value of the
+              nearest input pixel. This method is simple and fast but can result in blocky or pixelated outputs.
+            - ``"bilinear"``: Bilinear interpolation. Each output pixel is a weighted average of the four nearest input
+              pixels, computed using bilinear interpolation. This method produces smoother results compared
+              to nearest neighbor interpolation.
+            - ``"bilinear_v2"``: The optimized variant of
+              ``"bilinear"``, it may achieve better result(higher precision and speed) in some cases.
+
         extrapolation_value (float, optional): An optional float value used extrapolation, if applicable.
             Default: ``0.0`` .
 
@@ -358,7 +365,6 @@ class CropAndResize(Primitive):
         self.method = method
         validator.check_value_type("extrapolation_value", extrapolation_value, [float], self.name)
         self.extrapolation_value = extrapolation_value
-        self.is_ge = context.get_context("enable_ge")
 
 
 class NonMaxSuppressionV3(Primitive):
@@ -470,7 +476,7 @@ class NonMaxSuppressionWithOverlaps(Primitive):
     Raises:
         TypeError: If the dtype of `overlaps` , `scores` `overlap_threshold` and `score_threshold`
                    is not float16, float32 or float64.
-        TypeError: If `overlaps` or `scores` is not Tensor。
+        TypeError: If `overlaps` or `scores` is not Tensor.
         TypeError: If `max_output_size` is not Tensor or Scalar.If `max_output_size` is not int32.
         TypeError: If `overlap_threshold` is not Tensor or scalar. If its type is not float16, float32 or float64.
         TypeError: If `score_threshold` is not Tensor or scalar. If its type is not float16, float32 or float64.
@@ -1115,13 +1121,13 @@ class CombinedNonMaxSuppression(Primitive):
               it exceeds `max_total_size`.
 
     Inputs:
-        - **boxes** (Tensor) - A float32 Tensor with shape :math:`(batch_size, num_boxes, q, 4)`
+        - **boxes** (Tensor) - A float32 Tensor with shape :math:`(batch\_size, num\_boxes, q, 4)`
           representing the bounding box coordinates.
           `q` indicates mapping relationship between boxes and classes.
           If `q` is 1, all classes use the same bounding box. If `q` is equal to the number of classes,
           class-specific boxes are applied.
         - **scores** (Tensor) - A 3-D Tensor of float32 type with the shape
-          :math:`(batch_size, num_boxes, num_classes)`. It contains a score value for each box,
+          :math:`(batch\_size, num\_boxes, num\_classes)`. It contains a score value for each box,
           with each row of `boxes` represented by a single score.
         - **max_output_size_per_class** (Tensor) - The maximum number of boxes that can be selected for each class
           by the non-maximum suppression algorithm, represented by a scalar Tensor of type int32.

mindspore/ops/operations/inner_ops.py

@@ -238,13 +238,14 @@ class LambApplyOptimizerAssign(PrimitiveWithInfer):
     @prim_attr_register
     def __init__(self):
         """Initialize LambApplyOptimizerAssign"""
+        self.var_shape = "var_shape"
         self.add_prim_attr('side_effect_mem', True)
 
     def infer_shape(self, grad_shape, v_shape, m_shape, var_shape, beta1_shape, sub1_shape,
                     beta2_shape, sub2_shape, eps_shape, steps_shape, use_weight_shape, weight_decay_shape):
-        validator.check("var_shape", var_shape, "m_shape", m_shape, validator.EQ, self.name)
-        validator.check("var_shape", var_shape, "v_shape", v_shape, validator.EQ, self.name)
-        validator.check("var_shape", var_shape, "grad_shape", grad_shape, validator.EQ, self.name)
+        validator.check(self.var_shape, var_shape, "m_shape", m_shape, validator.EQ, self.name)
+        validator.check(self.var_shape, var_shape, "v_shape", v_shape, validator.EQ, self.name)
+        validator.check(self.var_shape, var_shape, "grad_shape", grad_shape, validator.EQ, self.name)
         return m_shape, v_shape, m_shape
 
     def infer_dtype(self, grad_dtype, v_dtype, m_dtype, var_dtype, beta1_dtype, sub1_dtype,
@@ -658,3 +659,25 @@ class ScaleGrad(PrimitiveWithInfer):
     @prim_attr_register
     def __init__(self):
         """Initialize ScaleGrad"""
+
+
+class KVCacheMgr(Primitive):
+    """
+    Update past with cur and index along sequence axis.
+
+    Inputs:
+        - **past** (Parameter) - 4-D tensor with shape: :math:`(batch_size, num_head, seq_len, hidden_size)`.
+        - **cur** (Tensor) - 4-D tensor with shape: :math:`(batch_size, num_head, 1, hidden_size)`.
+        - **index** (Tensor) - 1-D tensor with shape: :math:`(batch_size,)`.
+
+    Outputs:
+        Tensor, has the same data type and shape as original `past`.
+
+    Supported Platforms:
+        ``Ascend``
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        self.init_prim_io_names(inputs=['past', 'cur', 'index'], outputs=['past'])
+        self.add_prim_attr('side_effect_mem', True)