mindspore 2.4.10__cp311-cp311-win_amd64.whl → 2.6.0rc1__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

@@ -80,6 +80,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 +126,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"