mindspore 2.4.10__cp311-cp311-win_amd64.whl → 2.6.0__cp311-cp311-win_amd64.whl

mindspore/common/auto_dynamic_shape.py

@@ -17,15 +17,16 @@
 """Providing auto dynamic shape interface methods."""
 
 import os
+import math
 from mindspore import log as logger
-from mindspore._c_expression import GraphExecutor_, Tensor
+from mindspore._c_expression import GraphExecutor_, JitExecutor_, TensorPy as Tensor
 from mindspore.common._utils import is_shape_unknown, is_dim_unknown
 from mindspore.common.parameter import Parameter
 
 SHAPE_DIM_ANY = -1
 SHAPE_RANK_ANY = -2
 
-auto_dynamic_shepe_dict = {}
+auto_dynamic_shape_dict = {}
 
 
 class _AutoDynamicShapeManager:
@@ -37,8 +38,6 @@ class _AutoDynamicShapeManager:
         self.generalize_shape_cache = []
         self.real_phase_and_compile_args_dict = {}
         self.generalize_phase_and_compile_args_dict = {}
-        self._graph_executor = GraphExecutor_.get_instance()
-
 
     def __del__(self):
         self.real_shape_cache = []
@@ -167,7 +166,8 @@ class _AutoDynamicShapeManager:
             else:
                 delete_cache = set()
                 delete_cache.add(delete_phase)
-                self._graph_executor.del_net_res(None, delete_cache)
+                GraphExecutor_.get_instance().del_net_res(None, delete_cache)
+                JitExecutor_.get_instance().del_net_res(None, delete_cache)
                 del self.generalize_phase_and_compile_args_dict[delete_phase]
 
             # step3: delete compile args
@@ -186,8 +186,8 @@ class _AutoDynamicShapeManager:
             if isinstance(arg, Tensor) and isinstance(cache, Tensor):
                 if not self.is_tensor_equal(arg, cache):
                     return False
-            elif isinstance(arg, int) and isinstance(cache, int):
-                if arg != cache:
+            elif isinstance(arg, (int, float)) and isinstance(cache, (int, float)):
+                if not math.isclose(arg, cache):
                     return False
             elif isinstance(arg, (tuple, list)) and isinstance(cache, (tuple, list)):
                 if not self._compare_input_args_and_cache_args(arg, cache):
@@ -199,22 +199,19 @@ class _AutoIdentifyDynamicShape:
     """
     Represents a function auto identify dynamic shape.
     """
-    def __init__(self):
+    def __init__(self, auto_dynamic=False):
         self.all_shape_cache = {}
         self.is_sink_mode = False
-        self.is_enable_auto_dynamic_shape = True
+        self.is_enable_auto_dynamic_shape = auto_dynamic or os.getenv('MS_AUTO_DYNAMIC_SHAPE_ENABLE') == '1'
         self.save_cache_number = 3
-        self.enable_auto_identify = os.getenv('MS_AUTO_DYNAMIC_SHAPE_ENABLE')
         self.auto_dynamic_shape_manager = _AutoDynamicShapeManager()
 
-
     def __del__(self):
         self.all_shape_cache = {}
         self.is_sink_mode = False
         self.is_enable_auto_dynamic_shape = True
         self.save_cache_number = 3
 
-
     def _check_input_args_number(self, args_list):
         """check input arg number"""
         if self.auto_dynamic_shape_manager.get_real_shape_cache_number() > 0:
@@ -223,6 +220,12 @@ class _AutoIdentifyDynamicShape:
                 return False
         return True
 
+    def _check_input_int_float(self, arg, cache):
+        """check arg type is int or float"""
+        if (isinstance(arg, int) and not isinstance(cache, int)) or \
+           (isinstance(arg, float) and not isinstance(cache, float)):
+            return False
+        return math.isclose(arg, cache)
 
     def _check_input_tensor_type(self, args_list, cache_list):
         """check input args type"""
@@ -236,19 +239,15 @@ class _AutoIdentifyDynamicShape:
                 res = self._check_input_tensor_type(arg, cache)
                 if not res:
                     return False
-            elif (isinstance(arg, int) and isinstance(cache, int)) or \
-                 (isinstance(arg, float) and isinstance(cache, float)):
-                if arg != cache:
+            elif isinstance(arg, (int, float)):
+                if not self._check_input_int_float(arg, cache):
                     return False
             elif isinstance(arg, Tensor) and not isinstance(cache, Tensor):
                 return False
-            elif isinstance(arg, (int, float)) and not isinstance(cache, (int, float)):
-                return False
             elif isinstance(arg, (tuple, list)) and not isinstance(cache, (tuple, list)):
                 return False
         return True
 
-
     def _check_input_number_and_type(self, args_list):
         """check input number and type"""
         res = self._check_input_args_number(args_list)
@@ -262,18 +261,10 @@ class _AutoIdentifyDynamicShape:
                 return False
         return True
 
-
     def _is_enable_auto_dynamic_shape(self, args_list, is_sink_mode):
         """is enable auto identify shape"""
         if not is_sink_mode and not args_list:
             return False
-
-        if not self.enable_auto_identify:
-            self.enable_auto_identify = "0"
-
-        if self.enable_auto_identify == "0":
-            return False
-
         for elem in args_list:
             if elem is None:
                 continue
@@ -285,7 +276,6 @@ class _AutoIdentifyDynamicShape:
                 return self._is_enable_auto_dynamic_shape(elem, is_sink_mode)
         return True
 
-
     @staticmethod
     def _do_generalize_in_sink(input_arg, cache, input_index, cache_index, cache_type):
         """do generalize in sink, input rank must be 2"""
@@ -340,8 +330,9 @@ class _AutoIdentifyDynamicShape:
 
     def auto_dynamic_generate_compile_args(self, args_list, is_sink_mode):
         """generate compile args in auto dynamic shape"""
-        if not self._check_input_number_and_type(args_list) or \
-            not self._is_enable_auto_dynamic_shape(args_list, is_sink_mode):
+        if not self.is_enable_auto_dynamic_shape or \
+            not self._is_enable_auto_dynamic_shape(args_list, is_sink_mode) or \
+            not self._check_input_number_and_type(args_list):
             self.is_enable_auto_dynamic_shape = False
             return args_list
         self.is_sink_mode = is_sink_mode
@@ -415,11 +406,12 @@ class _AutoIdentifyDynamicShape:
             return Tensor(shape=shape_tuple, dtype=arg.dtype)
 
         def generalize_sequence(arg, cache, i):
+            if not hasattr(arg, "__ms_mutable__"):
+                return arg, False
             if is_sink_mode:
                 # when is_sink_mode=True, input must be the shape of Tensor.
                 res = self._do_generalize_in_sink(arg, cache, i, index, cache_type)
                 return res
-
             res = self._do_generalize_one_input_shape(arg, cache, cache_type, index, is_sink_mode)
             return res
 
@@ -437,12 +429,16 @@ class _AutoIdentifyDynamicShape:
                 res = generalize_tensor(arg, cache, i)
                 generalize_one_shape.append(res)
             elif isinstance(arg, (tuple, list)) and isinstance(cache, (tuple, list)):
-                res = generalize_sequence(arg, cache, i)
-                generalize_one_shape.append(res)
-            elif isinstance(arg, int) and isinstance(cache, int):
+                gen_seq_shape, can_generalize = generalize_sequence(arg, cache, i)
+                if not can_generalize:
+                    return input_args, False
+                if isinstance(arg, tuple):
+                    gen_seq_shape = tuple(gen_seq_shape)
+                generalize_one_shape.append(gen_seq_shape)
+            elif isinstance(arg, (int, float)) and isinstance(cache, (int, float)):
                 # when is_sink_mode=False, the input must may be scalar, or the value of list/tuple.
                 # is_sink_mode can not be True
-                if arg == cache:
+                if math.isclose(arg, cache):
                     generalize_one_shape.append(arg)
                 else:
                     logger.info("In auto dynamic shape mode, scalar/tuple/list must be equal, it can not be " \
@@ -479,29 +475,30 @@ class _AutoIdentifyDynamicShape:
 _auto_dynamic_shape = _AutoIdentifyDynamicShape()
 
 
-def get_auto_dynamic_shape_args(compile_args, key_id):
+def get_auto_dynamic_shape_args(compile_args, key_id, enable_auto_dynamic=False):
     """get auto dynamic shape args."""
-    if key_id not in auto_dynamic_shepe_dict:
-        auto_dynamic_shepe_dict[key_id] = _AutoIdentifyDynamicShape()
-    compile_args = auto_dynamic_shepe_dict[key_id].auto_dynamic_generate_compile_args(compile_args, False)
+    if key_id not in auto_dynamic_shape_dict:
+        auto_dynamic_shape_dict[key_id] = _AutoIdentifyDynamicShape(enable_auto_dynamic)
+    compile_args = auto_dynamic_shape_dict[key_id].auto_dynamic_generate_compile_args(compile_args, False)
     return compile_args
 
 
 def update_auto_dynamic_shape_phase(compile_args, key_id, phase):
     """update auto dynamic shape phase."""
-    if key_id in auto_dynamic_shepe_dict:
-        auto_dynamic_shepe_dict[key_id].update_phase_and_compile_args(compile_args, phase, False)
+    if key_id in auto_dynamic_shape_dict:
+        auto_dynamic_shape_dict[key_id].update_phase_and_compile_args(compile_args, phase, False)
 
 
-def get_auto_dynamic_shape_args_with_check_input_signature(compile_args, key_id, input_signature):
+def get_auto_dynamic_shape_args_with_check_input_signature(compile_args, key_id, input_signature,
+                                                           enable_auto_dynamic=False):
     """get auto dynamic shape args."""
     if input_signature is None:
-        return get_auto_dynamic_shape_args(compile_args, key_id)
+        return get_auto_dynamic_shape_args(compile_args, key_id, enable_auto_dynamic)
     return compile_args
 
 
 def update_auto_dynamic_shape_phase_with_check_input_signature(compile_args, key_id, phase, input_signature):
     """update auto dynamic shape phase."""
     if input_signature is None:
-        if key_id in auto_dynamic_shepe_dict:
-            auto_dynamic_shepe_dict[key_id].update_phase_and_compile_args(compile_args, phase, False)
+        if key_id in auto_dynamic_shape_dict:
+            auto_dynamic_shape_dict[key_id].update_phase_and_compile_args(compile_args, phase, False)

mindspore/common/dtype.py

@@ -59,58 +59,58 @@ __all__.extend(__dtype__)
 __all__.extend(__method__)
 
 # type definition
-bool_ = typing.Bool()
+bool_ = typing.kBool
 
-qint4x2 = typing.Int(4)
-int8 = typing.Int(8)
+qint4x2 = typing.kInt4
+int8 = typing.kInt8
 byte = int8
-int16 = typing.Int(16)
+int16 = typing.kInt16
 short = int16
-int32 = typing.Int(32)
+int32 = typing.kInt32
 intc = int32
-int64 = typing.Int(64)
+int64 = typing.kInt64
 intp = int64
 
-uint8 = typing.UInt(8)
+uint8 = typing.kUInt8
 ubyte = uint8
-uint16 = typing.UInt(16)
+uint16 = typing.kUInt16
 ushort = uint16
-uint32 = typing.UInt(32)
+uint32 = typing.kUInt32
 uintc = uint32
-uint64 = typing.UInt(64)
+uint64 = typing.kUInt64
 uintp = uint64
 
-float16 = typing.Float(16)
+float16 = typing.kFloat16
 half = float16
-float32 = typing.Float(32)
+float32 = typing.kFloat32
 single = float32
-float64 = typing.Float(64)
+float64 = typing.kFloat64
 double = float64
-bfloat16 = typing.BFloat(16)
-complex64 = typing.Complex(64)
-complex128 = typing.Complex(128)
-
-number = typing.Number()
-int_ = typing.Int()
-uint = typing.UInt()
-float_ = typing.Float()
-string = typing.String()
-list_ = typing.List()
-tuple_ = typing.Tuple()
-type_none = typing.TypeNone()
-_null = typing.TypeNull()
-
-tensor_type = typing.TensorType()
-index_slices = typing.RowTensorType()
-coo_tensor = typing.COOTensorType()
-csr_tensor = typing.CSRTensorType()
+bfloat16 = typing.kBFloat16
+complex64 = typing.kComplex64
+complex128 = typing.kComplex128
+
+number = typing.kNumber
+int_ = typing.kInt
+uint = typing.kUInt
+float_ = typing.kFloat
+string = typing.kString
+list_ = typing.kList
+tuple_ = typing.kTuple
+type_none = typing.kTypeNone
+_null = typing.kTypeNull
+
+tensor_type = typing.kTensorType
+index_slices = typing.kRowTensorType
+coo_tensor = typing.kCOOTensorType
+csr_tensor = typing.kCSRTensorType
 undetermined = typing.UndeterminedType()
 
 function = typing.Function()
 symbolic_key = typing.SymbolicKeyType()
-env_type = typing.EnvType()
-type_type = typing.TypeType()
-type_refkey = typing.RefKeyType()
+env_type = typing.kTypeEnv
+type_type = typing.kTypeType
+type_refkey = typing.kRefKeyType
 
 Int = typing.Int
 Float = typing.Float
@@ -279,8 +279,11 @@ def dtype_to_nptype(type_):
         complex64: np.complex64,
         complex128: np.complex128,
     }
-    if np_version_valid(False):
-        _dtype_nptype_dict.update({bfloat16: np_bfloat16})
+    if type_ == bfloat16:
+        if not np_version_valid(True):
+            raise TypeError("The Numpy bfloat16 data type is not supported now, please ensure that the current "
+                            "Numpy version is not less than the version when the mindspore is compiled.")
+        return np_bfloat16
     return _dtype_nptype_dict[type_]
 
 

mindspore/common/dump.py

@@ -27,17 +27,19 @@ def set_dump(target, enabled=True):
     `target` should be an instance of :class:`mindspore.nn.Cell` or :class:`mindspore.ops.Primitive` .
     Please note that this API takes effect only when Synchronous Dump is enabled and the `dump_mode`
     field in dump config file is ``"2"`` . See the `dump document
-    <https://www.mindspore.cn/docs/en/master/model_train/debug/dump.html>`_ for details.
+    <https://www.mindspore.cn/tutorials/en/master/debug/dump.html>`_ for details.
     The default enabled status for
     a :class:`mindspore.nn.Cell` or :class:`mindspore.ops.Primitive` is False.
 
     Note:
-        1. This API only supports being called before training starts.
+        1. This API is only effective for GRAPH_MODE whose graph compilation level is O0/O1 with Ascend backend,
+           and can not work for fusion Primitive operators.
+        2. This API only supports being called before training starts.
            If you call this API during training, it may not be effective.
-        2. After using `set_dump(Cell, True)` , operators in forward and backward
+        3. After using `set_dump(Cell, True)` , operators in forward and backward
            computation  (computation generated by the grad operations) of the
            cell will be dumped.
-        3. For :class:`mindspore.nn.SoftmaxCrossEntropyWithLogits` layer, the forward
+        4. For :class:`mindspore.nn.SoftmaxCrossEntropyWithLogits` layer, the forward
            computation and backward computation use the same set of
            operators. So you can only see dump data from backward computation.
            Please note that :class:`mindspore.nn.SoftmaxCrossEntropyWithLogits` layer will also use
@@ -57,14 +59,15 @@ def set_dump(target, enabled=True):
         .. note::
             Please set environment variable `MINDSPORE_DUMP_CONFIG` to the dump config file and set `dump_mode` field
             in dump config file to 2 before running this example.
-            See `dump document <https://www.mindspore.cn/docs/en/master/model_train/debug/dump.html>`_ for details.
+            See `dump document <https://www.mindspore.cn/tutorials/en/master/debug/dump.html>`_ for details.
 
         >>> import numpy as np
         >>> import mindspore as ms
         >>> import mindspore.nn as nn
         >>> from mindspore import Tensor, set_dump
         >>>
-        >>> ms.set_context(device_target="Ascend", mode=ms.GRAPH_MODE)
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
+        >>> ms.set_device(device_target="Ascend")
         >>>
         >>> class MyNet(nn.Cell):
         ...     def __init__(self):

mindspore/common/file_system.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ============================================================================
 """File system registration management"""
+from mindspore import log as logger
 
 
 class FileSystem:
@@ -40,11 +41,18 @@ def _register_mindio_file_system(fs: FileSystem):
         import mindio
     except ImportError:
         return False
-    if mindio.initialize() != 0:
+    try:
+        ret = mindio.initialize()
+    except AttributeError as e:
+        logger.warning(f"Failed to initialize MindIO: {e}")
+        return False
+    if ret != 0:
+        logger.warning(f"Failed to initialize MindIO: ret = {ret}")
         return False
     fs.create = mindio.create_file
     fs.create_args = ()
     fs.open = mindio.open_file
     fs.open_args = ()
     fs.backend = "mindio"
+    logger.info("The weights are stored using MindIO as the backend.")
     return True

mindspore/common/generator.py

@@ -56,6 +56,9 @@ class Generator:
     A generator that manages the state of random numbers and provides seed and offset for random functions.
     When the seed and offset are fixed, the random function generates the same random sequence.
 
+    .. note::
+        Graph mode does not support the use of multiple generators at the same time for now.
+
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
 
@@ -80,6 +83,7 @@ class Generator:
             Tensor(0, mstype.int64), name="offset", requires_grad=False)
 
         self._generator = GeneratorOp().set_device("CPU")
+        self._generator.add_prim_attr("manual_seed", False)
         self._to_scalar = TensorToScalar()
 
     def set_state(self, state):
@@ -125,6 +129,7 @@ class Generator:
             raise TypeError("Seed must be an integer.")
         seed = Tensor(seed, mstype.int64)
         self._generator(MANUAL_SEED, (self._seed, self._offset, seed))
+        self._generator.add_prim_attr("manual_seed", True)
         return self
 
     def initial_seed(self):

mindspore/common/hook_handle.py

@@ -15,7 +15,7 @@
 """The removable handle for cell hook function."""
 from __future__ import absolute_import
 import weakref
-from mindspore._c_expression import Tensor as Tensor_
+from mindspore._c_expression import TensorPy as Tensor_
 
 
 class _TensorHookHandle:
@@ -29,8 +29,9 @@ class _TensorHookHandle:
         ``Ascend`` ``GPU`` ``CPU``
     """
 
-    def __init__(self):
+    def __init__(self, tensor):
         self.id = None
+        self.tensor_ref = weakref.ref(tensor)
 
     def remove(self):
         """
@@ -66,6 +67,9 @@ class _TensorHookHandle:
         """
         if self.id is not None:
             Tensor_.remove_hook(self.id)
+            tensor = self.tensor_ref()
+            if tensor is not None:
+                tensor._remove_hook() # pylint:disable=protected-access
 
 
 class HookHandle:

mindspore/common/initializer.py

@@ -36,9 +36,7 @@ class Initializer:
     Note:
         Initializers are intended to be used for delayed initialization in parallel mode rather than Tensor
         initialization. If you have to use Initializers to create a Tensor, :func:`mindspore.Tensor.init_data` should be
-        followed in most of the cases. For more information, please refer to `mindspore.Tensor.init_data
-        <https://www.mindspore.cn/docs/en/master/api_python/mindspore/Tensor/mindspore.Tensor.init_data.html#
-        mindspore-tensor-init-data>`_ .
+        followed in most of the cases. For more information, please refer to :func:`mindspore.Tensor.init_data` .
 
     Args:
         kwargs (dict): Keyword arguments for Initializer.
@@ -371,9 +369,9 @@ class HeUniform(Initializer):
     .. math::
         boundary = \text{gain} \times \sqrt{\frac{3}{fan\_mode}}
 
-    where :math:`gain` is an optional scaling factor. If :math:`fan\_mode` is ``'fan_in'``,
-    it is the number of input units of the weight tensor. If :math:`fan\_mode` is ``'fan_out'``,
-    it is the number of output units of the weight tensor.
+    where :math:`gain` is an optional scaling factor. If `mode` is ``'fan_in'``,
+    :math:`fan\_mode` in the formula is the number of input units of the weight tensor. If `mode` is ``'fan_out'``,
+    :math:`fan\_mode` is the number of output units of the weight tensor.
 
     For details of HeUniform algorithm, please check
     `<https://arxiv.org/abs/1502.01852>`_.
@@ -653,12 +651,16 @@ class Orthogonal(Initializer):
 class VarianceScaling(Initializer):
     r"""
     Generates an random array with scaling in order to initialize a tensor.
+
     When `distribution` is 'truncated_normal' or 'untruncated_normal', the value will be sampled from truncated or
     untruncated normal distribution with a mean of 0 and a scaled standard deviation
-    :math:`stddev = \sqrt{\frac{scale}{n}}`. :math:`n` will be the number of input units if `mode` is ``'fan_in'``,
+    :math:`stddev = \sqrt{\frac{scale}{n}}`.
+
+    :math:`n` will be the number of input units if `mode` is ``'fan_in'``,
     while :math:`n` will be
     the number of output units if `mode` is ``'fan_out'``. :math:`n` will be the average of ``'fan_in'``
     and ``'fan_out'`` if `mode` is ``'fan_avg'``.
+
     When `distribution` is ``'uniform'``, the value will be sampled from a uniform distribution within the limit of
     :math:`[-\sqrt{\frac{3*scale}{n}}, \sqrt{\frac{3*scale}{n}}]`.
 
@@ -827,6 +829,8 @@ def initializer(init, shape=None, dtype=mstype.float32):
     Args:
         init (Union[Tensor, str, Initializer, numbers.Number]): Initialize value.
 
+            - `Tensor`: The tensor will be called to initialize tensor.
+
             - `str`: The `init` should be the alias of the class inheriting from `Initializer` and the corresponding
               class will be called in practice. The value of `init` can be ``"normal"``, ``"ones"`` or
               ``"zeros"``, etc.
@@ -835,13 +839,12 @@ def initializer(init, shape=None, dtype=mstype.float32):
 
             - `numbers.Number`: The `Constant` will be called to initialize tensor.
 
-            - `Tensor`: The tensor will be called to initialize tensor.
-
         shape (Union[tuple, list, int]): The shape of the initialized tensor. Default: ``None`` .
         dtype (:class:`mindspore.dtype`): The type of data in initialized tensor. Default: ``mstype.float32`` .
 
     Returns:
-        Tensor, return is Tensor object.
+        Returns a Tensor with the shape specified by the input `shape`. If `shape` is ``None``,
+        the returned Tensor will have the same shape as ``init``.
 
     Raises:
         TypeError: The type of the argument 'init' is not correct.

mindspore/common/jit_begin_end.py

@@ -0,0 +1,94 @@
+# Copyright 2024 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.
+# ============================================================================
+
+"""JIT begin/end for various JIT context compile."""
+
+from .jit_trace import _jit_trace_begin, _jit_trace_end
+
+
+def _jit_begin(fn_name, *args):
+    """
+    Start to build a MindIR func graph for a code snippet.
+
+    This allows the MindSpore runtime to apply optimizations based on generated func graph.
+
+    Note:
+        Use it with `jit_end` cooperatively.
+
+    Also see: :func:`jit_end`.
+
+    Args:
+        fn_name (str): The name of func graph to be built.
+        args (tuple): The arguments of func graph.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> from mindspore import Tensor
+        >>> from mindspore.common.jit_begin_end import _jit_begin as jit_begin
+        >>> from mindspore.common.jit_begin_end import _jit_end as jit_end
+        ...
+        >>> x = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
+        >>> y = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
+        >>> def tensor_add(x, y):
+        ...     jit_begin(x, y)
+        ...     z = x + y
+        ...     z = jit_end(z)
+        ...     return z
+        ...
+        >>> out = tensor_add(x, y)
+    """
+    return _jit_trace_begin(fn_name, *args)
+
+
+def _jit_end(*output_args):
+    """
+    Finish building a MindIR func graph for a code snippet.
+
+    This allows the MindSpore runtime to apply optimizations based on generated func graph.
+
+    Note:
+        Use it with `jit_begin` cooperatively.
+
+    Also see: :func:`jit_begin`.
+
+    Args:
+        output_args (tuple): The output of func graph.
+
+    Returns:
+        The same as args `output_args`.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> from mindspore import Tensor
+        >>> from mindspore.common.jit_begin_end import _jit_begin as jit_begin
+        >>> from mindspore.common.jit_begin_end import _jit_end as jit_end
+        ...
+        >>> x = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
+        >>> y = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
+        >>> def tensor_add(x, y):
+        ...     jit_begin(x, y)
+        ...     z = x + y
+        ...     z = jit_end(z)
+        ...     return z
+        ...
+        >>> out = tensor_add(x, y)
+    """
+    return _jit_trace_end(*output_args)

mindspore/common/jit_config.py

@@ -21,7 +21,7 @@ class JitConfig:
     Args:
         jit_level (str, optional): Used to control the compilation optimization level.
             Supports ["O0", "O1", "O2"]. Default: ``""`` , The framework automatically selects the execution method.
-            Not recommended, it is recommended to use the jit decorator.
+            Not recommended, it is recommended to use the JIT decorator.
 
             - ``"O0"``: Except for optimizations that may affect functionality, all other optimizations are turned off,
               adopt KernelByKernel execution mode.
@@ -96,3 +96,8 @@ class JitConfig:
         self.jit_config_dict["jit_syntax_level"] = jit_syntax_level
         self.jit_config_dict["debug_level"] = debug_level
         self.jit_config_dict["infer_boost"] = infer_boost
+        if "backend" not in self.jit_config_dict:
+            if jit_level == "O0" or jit_level == "O1":
+                self.jit_config_dict["backend"] = "ms_backend"
+            elif jit_level == "O2":
+                self.jit_config_dict["backend"] = "GE"