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

mindspore/common/file_system.py

@@ -14,10 +14,14 @@
 # ============================================================================
 """File system registration management"""
 from mindspore import log as logger
+from mindspore import _checkparam as Validator
+
+mindio_server_info = {"memfs.data_block_pool_capacity_in_gb": "100"}
 
 
 class FileSystem:
     """File operation interface manager"""
+
     def __init__(self):
         self.create = open
         self.create_args = ("ab",)
@@ -35,20 +39,33 @@ def _register_basic_file_system(fs: FileSystem):
     return True
 
 
-def _register_mindio_file_system(fs: FileSystem):
-    """register mindio file system"""
+def _init_mindio():
+    """Initialize MindIO and return the module if successful"""
     try:
-        import mindio
+        import mindio_acp as mindio
+        ret = mindio.initialize(server_info=mindio_server_info)
+        if ret == 0:
+            return mindio
+        logger.warning(f"Failed to initialize mindio_acp: ret = {ret}")
     except ImportError:
-        return False
+        pass
     try:
+        import mindio
         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}")
+        if ret == 0:
+            return mindio
+        logger.warning(f"Failed to initialize mindio: ret = {ret}")
+    except ImportError:
+        pass
+    return None
+
+
+def _register_mindio_file_system(fs: FileSystem):
+    """register mindio file system"""
+    mindio = _init_mindio()
+    if mindio is None:
         return False
+
     fs.create = mindio.create_file
     fs.create_args = ()
     fs.open = mindio.open_file
@@ -56,3 +73,36 @@ def _register_mindio_file_system(fs: FileSystem):
     fs.backend = "mindio"
     logger.info("The weights are stored using MindIO as the backend.")
     return True
+
+
+def set_mindio_server_info(data_block_pool_capacity_in_gb=100):
+    """
+    Configure MindIO server settings.
+
+    Args:
+        data_block_pool_capacity_in_gb (int): Memory pool capacity for data blocks in gigabytes.
+    """
+    global mindio_server_info
+    Validator.check_positive_int(data_block_pool_capacity_in_gb, "data_block_pool_capacity_in_gb")
+    mindio_server_info["memfs.data_block_pool_capacity_in_gb"] = str(data_block_pool_capacity_in_gb)
+
+
+def mindio_preload(ckpt_file_name):
+    """
+    Preload data into memory using MindIO for faster access.
+
+    Args:
+        ckpt_file_name (str): Checkpoint file name.
+
+    Returns:
+        bool: True if preloading is successful, False otherwise.
+    """
+    Validator.check_value_type('ckpt_file_name', ckpt_file_name, str, "mindio_preload")
+    mindio = _init_mindio()
+    if mindio is None:
+        return False
+    if not hasattr(mindio, 'preload'):
+        logger.warning("MindIO module does not have preload method")
+        return False
+    mindio.preload(ckpt_file_name)
+    return True

mindspore/common/generator.py

@@ -77,10 +77,9 @@ class Generator:
     """
 
     def __init__(self):
-        self._seed = Parameter(Tensor(0, mstype.int64),
-                               name="seed", requires_grad=False)
+        self._seed = Parameter(Tensor(0, mstype.int64), requires_grad=False)
         self._offset = Parameter(
-            Tensor(0, mstype.int64), name="offset", requires_grad=False)
+            Tensor(0, mstype.int64), requires_grad=False)
 
         self._generator = GeneratorOp().set_device("CPU")
         self._generator.add_prim_attr("manual_seed", False)

mindspore/common/hook_handle.py

@@ -16,6 +16,21 @@
 from __future__ import absolute_import
 import weakref
 from mindspore._c_expression import TensorPy as Tensor_
+from mindspore._check_jit_forbidden_api import jit_forbidden_register
+
+
+# Global variable to mark the `Parameter` hook and `Cell` hook version
+_HOOK_VERSION = 0
+
+
+def _update_hook_version():
+    global _HOOK_VERSION
+    _HOOK_VERSION += 1
+
+
+def _hook_version():
+    global _HOOK_VERSION
+    return _HOOK_VERSION
 
 
 class _TensorHookHandle:
@@ -31,8 +46,9 @@ class _TensorHookHandle:
 
     def __init__(self, tensor):
         self.id = None
-        self.tensor_ref = weakref.ref(tensor)
+        self.tensor_weakref = weakref.ref(tensor)
 
+    @jit_forbidden_register
     def remove(self):
         """
         Remove the tensor hook function, which corresponds to this '_TensorHookHandle' object.
@@ -67,9 +83,9 @@ class _TensorHookHandle:
         """
         if self.id is not None:
             Tensor_.remove_hook(self.id)
-            tensor = self.tensor_ref()
+            tensor = self.tensor_weakref()
             if tensor is not None:
-                tensor._remove_hook() # pylint:disable=protected-access
+                tensor._remove_hook()  # pylint:disable=protected-access
 
 
 class HookHandle:
@@ -81,20 +97,25 @@ class HookHandle:
         It is only supported in pynative mode and works when registering or removing hook function for Cell object.
 
     Args:
-        hook_dict (Dict): The hook object with hook function registered on. Default value: None.
+        hook_dict (Dict, optional): The hook object with hook function registered on. Default value: ``None`` .
+        extra_dict (Dict, optional): The extra dict. Default value: ``None`` .
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
     """
     unique_id = 0
 
-    def __init__(self, hook_dict=None):
+    def __init__(self, hook_dict=None, *, extra_dict=None):
         self.hook_dict_ref = None
+        self.extra_dict_ref = None
         if hook_dict is not None:
             self.hook_dict_ref = weakref.ref(hook_dict)
             self.handle_id = HookHandle.unique_id
             HookHandle.unique_id += 1
+            if extra_dict is not None:
+                self.extra_dict_ref = weakref.ref(extra_dict)
 
+    @jit_forbidden_register
     def remove(self):
         """
         Remove the cell hook function, which corresponds to this 'HookHandle' object.
@@ -141,7 +162,14 @@ class HookHandle:
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32,
             value= [ 2.00000000e+00]))
         """
+        _update_hook_version()  # pylint:disable=protected-access
+
         if self.hook_dict_ref is not None:
             hook_dict = self.hook_dict_ref()
             if hook_dict is not None and self.handle_id in hook_dict:
                 del hook_dict[self.handle_id]
+
+        if self.extra_dict_ref is not None:
+            extra_dict = self.extra_dict_ref()
+            if extra_dict is not None and self.handle_id in extra_dict:
+                del extra_dict[self.handle_id]

mindspore/common/jit_config.py

@@ -97,7 +97,7 @@ class JitConfig:
         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":
+            if jit_level in ["O0", "O1"]:
                 self.jit_config_dict["backend"] = "ms_backend"
             elif jit_level == "O2":
                 self.jit_config_dict["backend"] = "GE"

mindspore/common/jit_trace.py

@@ -17,16 +17,15 @@
 
 import inspect
 import re
+import types
 from functools import wraps
 import mindspore as ms
 from mindspore import log as logger
 from mindspore import context
 from mindspore.common.jit_context import JitContext, set_jit_context, jit_context
 from mindspore.common.tensor import Tensor as PythonTensor
-from mindspore._checkparam import is_stub_tensor
 from mindspore._c_expression import TraceRecorder as tr
 from mindspore._c_expression import JitExecutor_
-from mindspore._c_expression import TensorNode
 from mindspore._c_expression import TensorPy as Tensor, CSRTensor, COOTensor
 from mindspore._c_expression import typing
 
@@ -35,7 +34,7 @@ class TraceJitContext(JitContext):
     """JIT Context for trace JIT."""
 
     def __init__(self):
-        JitContext.__init__(self)
+        super().__init__()
         self._is_nested = False
 
     def set_is_nested(self, status):
@@ -44,29 +43,34 @@ class TraceJitContext(JitContext):
     def is_nested(self):
         return self._is_nested
 
+    def args_preprocess(self, prim_name, prim_res, *args):
+        args = tuple(_convert_arg_for_operators(arg, prim_name)
+                     for arg in args)
+        file_names, linenos = _get_caller_lines()
+        return prim_res, file_names, linenos, args
+
     def run_op(self, prim, prim_res, *args):
         """Capture op"""
         logger.debug(f'prim: {prim}, args: {args}, prim_res: {prim_res}')
-        if isinstance(prim_res, TensorNode):
-            prim_res = prim_res.get_value()
-        prim_res = _sync_stub_tensor(prim_res)
-        args = tuple(_sync_stub_tensor(arg) for arg in args)
-        args = tuple(_convert_arg_for_operators(arg, prim.name) for arg in args)
-        file_names, linenos = _get_caller_lines()
-        tr.get_instance().new_node(prim, prim_res, file_names, linenos, False, *args)
+        prim_res, file_names, linenos, args = self.args_preprocess(prim.name, prim_res, *args)
+        tr.get_instance().new_node(prim, (prim_res, file_names, linenos, False), *args)
         return prim_res
 
+    def prepare_op(self, prim_name, prim_res, *args):
+        """Prepare op"""
+        logger.debug(f'prim: {prim_name}, args: {args}, prim_res: {prim_res}')
+        return self.args_preprocess(prim_name, prim_res, *args)
+
     def run_graph(self, phase, prim_res, *args):
         """Capture func_graph generated from ast"""
         logger.debug(f'phase: {phase}, args: {args}, prim_res: {prim_res}')
-        if isinstance(prim_res, TensorNode):
-            prim_res = prim_res.get_value()
-        prim_res = _sync_stub_tensor(prim_res)
-        args = tuple(_sync_stub_tensor(arg) for arg in args)
         file_names, linenos = _get_caller_lines()
-        tr.get_instance().new_fg_node((phase, prim_res, file_names, linenos, self._is_nested), *args)
+        tr.get_instance().new_fg_node((prim_res, file_names, linenos, phase, self._is_nested), *args)
         return prim_res
 
+    def default_output(self):
+        return PythonTensor(0)
+
 
 _compile_only = False
 _trace_jit_context = TraceJitContext()
@@ -80,20 +84,6 @@ def _set_compile_only(compile_only=True):
     _compile_only = compile_only
 
 
-def _sync_stub_tensor(stub):
-    """Synchronize stub tensor"""
-    if is_stub_tensor(stub):
-        real_tensor = stub.stub_sync()
-        logger.debug(f'Convert stub tensor, stub: [{type(stub)}] {id(stub)}/{stub}, '
-                     f'tensor: [{type(real_tensor)}] {id(real_tensor)}/{real_tensor}')
-        return real_tensor
-    if isinstance(stub, tuple):
-        return tuple(_sync_stub_tensor(item) for item in stub)
-    if isinstance(stub, list):
-        return list(_sync_stub_tensor(item) for item in stub)
-    return stub
-
-
 def convert_tensorpy(args):
     new_args = []
     for arg in args:
@@ -129,84 +119,73 @@ def nested_run(obj, cell, *args):
     if res is not tuple:
         res = (res,)
     file_names, linenos = _get_caller_lines()
-    res = _sync_stub_tensor(res)
     set_jit_context(None)
     return file_names, linenos, res
 
 
-def _jit_trace(fn):
-    """
-    Create a callable MindSpore graph from a Python function by trace method.
-
-    This allows the MindSpore runtime to apply optimizations based on traced func graph.
-
-    Args:
-        fn (Function): The Python function that will be run as a graph. Default: ``None`` .
-
-    Returns:
-        Function, if `fn` is not None, returns a callable function that will execute the compiled function; If `fn` is
-        None, returns a decorator and when this decorator invokes with a single `fn` argument, the callable function is
-        equal to the case when `fn` is not None.
-
-    Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-
-    Examples:
-        >>> import numpy as np
-        >>> from mindspore import Tensor
-        >>> from mindspore.common.jit_trace import _jit_trace as jit_trace
-        ...
-        >>> x = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
-        >>> y = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
-        ...
-        >>> # To create a callable MindSpore graph by calling decorator @jit_trace
-        >>> def tensor_add(x, y):
-        ...     z = x + y
-        ...     return z
-        ...
-        >>> tensor_add_graph = jit_trace(fn=tensor_add)
-        >>> out = tensor_add_graph(x, y)
-    """
-
-    @wraps(fn)
-    def jit_trace_wrap(*args, **kwargs):
-        # If a trace graph is already built, keep going without building a new trace graph.
-        if jit_context():
-            return fn(*args, **kwargs)
-        # Start trace process.
-        if kwargs:
-            bound_arguments = inspect.signature(fn).bind(*args, **kwargs)
-            bound_arguments.apply_defaults()
-            args = bound_arguments.args
-            kwargs = bound_arguments.kwargs
-        generate_name = fn.__module__
-        if args:
-            jit_args = args[1:] if hasattr(args[0], fn.__name__) else args
-            obj = args[0]
-            if hasattr(obj, fn.__name__):  # Add class name for Cell.
-                generate_name = generate_name + "." + obj.__class__.__name__
-        else:
-            jit_args = args
-        generate_name = generate_name + "." + fn.__name__ + "#" + str(id(fn))
-        # Add create time for Cell.
-        if args and hasattr(obj, fn.__name__):
-            generate_name = generate_name + '#created_' + str(args[0].create_time)
-        line_str = fn.__code__.co_filename + ":" + str(fn.__code__.co_firstlineno)
-        generate_name = generate_name + '#[' + line_str + ']'
-
-        new_compile = _jit_trace_begin(generate_name, *jit_args)
-        if new_compile:
-            fn_res = fn(*args, **kwargs)
-            logger.debug(f'fn: {fn}, fn_res: {fn_res}, line: {line_str}')
-            # Use fn's output to build func graph's output.
-            output = _jit_trace_end(fn_res)
-        else:
-            output = _jit_trace_end(None)  # Run with compilation.
-        logger.debug(f'output: {output}')
-        return output
-
-    jit_trace_wrap.__trace_func__ = True
-    return jit_trace_wrap
+def _jit_trace():
+    """Return the wrapped function for trace mode jit."""
+    def wrap_func(fn):
+        if hasattr(fn, "construct"):
+            if isinstance(fn, ms.nn.Cell):
+                # Bound the cell object to get the self arg.
+                return types.MethodType(_jit_trace()(fn.construct.__func__), fn)
+            if isinstance(fn, type) and issubclass(fn, ms.nn.Cell):
+                fn.construct = _jit_trace()(fn.construct)
+            return fn
+
+        if isinstance(fn, types.MethodType):
+            return types.MethodType(_jit_trace()(fn.__func__), fn.__self__)
+
+        if not isinstance(fn, types.FunctionType):
+            logger.warning(f"The fn should be function, method or cell instance/class, but got {fn}")
+            return fn
+
+        if hasattr(fn, "__wrapped_by_jit__"):
+            logger.warning(f"The fn {fn} should be wrapped by jit only once.")
+
+        @wraps(fn)
+        def jit_trace_wrap(*args, **kwargs):
+            # If a trace graph is already built, keep going without building a new trace graph.
+            if jit_context():
+                return fn(*args, **kwargs)
+            # Start trace process.
+            if kwargs:
+                bound_arguments = inspect.signature(fn).bind(*args, **kwargs)
+                bound_arguments.apply_defaults()
+                args = bound_arguments.args
+                kwargs = bound_arguments.kwargs
+            generate_name = fn.__module__
+            if args:
+                jit_args = args[1:] if hasattr(args[0], fn.__name__) else args
+                obj = args[0]
+                if hasattr(obj, fn.__name__):  # Add class name for Cell.
+                    generate_name = generate_name + "." + obj.__class__.__name__
+            else:
+                jit_args = args
+            generate_name = generate_name + "." + fn.__name__ + "#" + str(id(fn))
+            # Add create time for Cell.
+            if args and hasattr(obj, fn.__name__):
+                generate_name = generate_name + '#created_' + str(args[0].create_time)
+            line_str = fn.__code__.co_filename + ":" + str(fn.__code__.co_firstlineno)
+            generate_name = generate_name + '#[' + line_str + ']'
+
+            new_compile = _jit_trace_begin(generate_name, *jit_args)
+            if new_compile:
+                fn_res = fn(*args, **kwargs)
+                logger.debug(f'fn: {fn}, fn_res: {fn_res}, line: {line_str}')
+                # Use fn's output to build func graph's output.
+                output = _jit_trace_end(fn_res)
+            else:
+                output = _jit_trace_end(None)  # Run with compilation.
+            logger.debug(f'output: {output}')
+            return output
+
+        jit_trace_wrap.__trace_func__ = True
+        setattr(jit_trace_wrap, "__wrapped_by_jit__", True)
+        return jit_trace_wrap
+
+    return wrap_func
 
 
 def _get_caller_lines():
@@ -240,6 +219,8 @@ def _get_args_for_run(args):
             new_args.append(arg)
         elif isinstance(arg, dict) and hasattr(arg, "__ms_mutable__"):
             new_args.append(tuple(arg.values()))
+        elif isinstance(arg, (tuple, list)) and hasattr(arg, "__ms_mutable__"):
+            new_args.append(arg)
     return tuple(new_args)
 
 
@@ -284,7 +265,6 @@ def _jit_trace_begin(fn_name, *args):
     logger.debug(f'_jit_trace_begin, args: {args}')
     _trace_jit_context.set_is_nested(False)
     set_jit_context(_trace_jit_context)
-    args = tuple(_sync_stub_tensor(arg) for arg in args)
     for arg in args:
         logger.debug(f'_jit_trace_begin, arg: {arg}, {type(arg)}')
 
@@ -355,7 +335,6 @@ def _jit_trace_end(*output_args):
         logger.debug(f'jit trace result: {output}')
     else:
         logger.debug(f'output_args: {output_args}')
-        output_args = tuple(_sync_stub_tensor(arg) for arg in output_args)
         file_names, linenos = _get_caller_lines()
         tr.get_instance().end_graph(file_names, linenos, *output_args)
         if _compile_only:

mindspore/common/np_dtype.py

@@ -16,10 +16,10 @@
 # ============================================================================
 """Numpy data type for MindSpore."""
 
-from mindspore._c_expression.np_dtypes import np_version_valid
-if np_version_valid(True):
+from mindspore._c_expression.np_dtypes import np_dtype_valid
+if np_dtype_valid(True):
     from mindspore._c_expression.np_dtypes import bfloat16 # pylint: disable=unused-import
 
 __all__ = []
-if np_version_valid(False):
+if np_dtype_valid(False):
     __all__.extend(["bfloat16"])

mindspore/common/parameter.py

@@ -33,6 +33,7 @@ from mindspore import context
 from mindspore.common._utils import get_slice_num, get_slice_shape
 from mindspore.common.initializer import initializer
 from mindspore.common.tensor import Tensor, _TensorMeta
+from mindspore.common.hook_handle import _update_hook_version
 from mindspore import _checkparam as Validator
 from mindspore._check_jit_forbidden_api import jit_forbidden_register
 from mindspore._c_expression import TensorPy as Tensor_
@@ -49,20 +50,27 @@ import mindspore.common._monad as monad
 __all__ = ['Parameter', 'ParameterTuple']
 
 PARAMETER_NAME_DEFAULT = "Parameter"
+_GENERATED_PARAMETER_NAME_PREFIX = PARAMETER_NAME_DEFAULT + '#'
 PARAMETER_NAME_PREFIX_MAX_LEN = 1024
 
-# Global variable for parameter unique key.
-_GLOBAL_PARAMETER_KEY = -1
+_PARAMETER_NAME_ID = 0
+
+
+def _generate_parameter_name():
+    global _PARAMETER_NAME_ID
+    name = _GENERATED_PARAMETER_NAME_PREFIX + str(_PARAMETER_NAME_ID)
+    _PARAMETER_NAME_ID += 1
+    return name
+
+
+def _is_parameter_generated(param_name):
+    if not param_name or not isinstance(param_name, str):
+        return False
+    return param_name.startswith(_GENERATED_PARAMETER_NAME_PREFIX)
 
-# Global variable to mark the hook of parameter is updated
-_parameter_hook_updated = True
-def set_parameter_hook_updated(value):
-    global _parameter_hook_updated
-    _parameter_hook_updated = value
 
-def parameter_hook_updated():
-    global _parameter_hook_updated
-    return _parameter_hook_updated
+# Global variable for parameter unique key.
+_GLOBAL_PARAMETER_KEY = -1
 
 
 def _is_in_auto_parallel_mode():
@@ -72,7 +80,7 @@ def _is_in_auto_parallel_mode():
 
 def _is_parallel_mode():
     """ Whether is parallel mode """
-    if not _is_initialized() or context.get_context('mode') == context.PYNATIVE_MODE:
+    if not _is_initialized():
         return False
     if os.getenv("RUN_MODE") != "predict":
         return False
@@ -410,10 +418,6 @@ class Parameter(Tensor_):
                                "1. context.set_ps_context(enable_ps=True) \n"
                                "2. export MS_ROLE environment variable \n"
                                "Please refer to the official website for detailed usage.")
-
-        if context.get_context("mode") == context.PYNATIVE_MODE:
-            raise RuntimeError("Parameter server training is not supported in pynative mode currently."
-                               "Please switch to graph mode and retry.")
         self.is_param_ps = True
         self.init_in_server = init_in_server
         self.param_info.init_in_server = init_in_server
@@ -496,11 +500,11 @@ class Parameter(Tensor_):
                 the default value `PARAMETER_NAME_DEFAULT` is used.
         """
         if name_ is None:
-            name_ = PARAMETER_NAME_DEFAULT
+            name_ = _generate_parameter_name()
         elif isinstance(name_, str):
             name_ = name_.strip()
             if name_ == '':
-                name_ = PARAMETER_NAME_DEFAULT
+                name_ = _generate_parameter_name()
             if len(name_) > PARAMETER_NAME_PREFIX_MAX_LEN:
                 raise ValueError("The length of the '{}' name should be less than {}.".
                                  format(name_, PARAMETER_NAME_PREFIX_MAX_LEN))
@@ -904,13 +908,10 @@ class Parameter(Tensor_):
         incoming_tensor_is_init = isinstance(data, Tensor) and not data.has_init
         current_tensor_is_init = isinstance(self, Tensor) and not self.has_init
         if self.dtype != data.dtype:
-            if mstype.implicit_conversion_seq.get(self.dtype) < mstype.implicit_conversion_seq.get(data.dtype):
-                self._raise_type_error(data.dtype)
-            else:
-                from mindspore.ops import functional as F
-                if isinstance(data, Tensor) and data.init is not None:
-                    data.init_data()
-                data = F.cast(data, self.dtype)
+            from mindspore.ops import functional as F
+            if isinstance(data, Tensor) and data.init is not None:
+                data.init_data()
+            data = F.cast(data, self.dtype)
         if isinstance(data, Tensor) and data.has_init:
             # The parameter has been initialized, directly update by the data
             if current_tensor_is_init:
@@ -938,7 +939,6 @@ class Parameter(Tensor_):
             init_data_args += (slice_index, layout[2], layout[5])
         return init_data_args
 
-
     def init_data(self, layout=None, set_sliced=False):
         """
         Initialize the parameter's data.
@@ -1007,11 +1007,11 @@ class Parameter(Tensor_):
         For details, please refer to :func:`mindspore.Tensor.register_hook`.
         """
         handle = Tensor.register_hook(self, hook_fn)
-        set_parameter_hook_updated(True)
+        _update_hook_version()
         return handle
 
     def _remove_hook(self):
-        set_parameter_hook_updated(True)
+        _update_hook_version()
 
     def _offload(self):
         r"""
@@ -1030,7 +1030,6 @@ class Parameter(Tensor_):
         """
         return Tensor_._offload(self, True)
 
-
     def _load(self):
         r"""
         Load parameter to device.
@@ -1160,6 +1159,5 @@ class ParameterTuple(tuple):
                 _insert_accumu_init_info(x1.name, init_to_value(init))
         return ParameterTuple(new)
 
-
     def __parameter_tuple__(self):
         """For parse check."""

mindspore/common/recompute.py

@@ -18,15 +18,15 @@ from collections import OrderedDict
 from types import MethodType
 from mindspore import log as logger
 from mindspore.nn.cell import Cell
-from mindspore import context
 from mindspore.common.tensor import Tensor
 from mindspore import ops
 from mindspore.ops.composite import GradOperation
 from mindspore.common._register_for_recompute import recompute_registry
 from mindspore.common.api import _pynative_executor, _no_grad
 from mindspore.common.generator import get_rng_state, set_rng_state
-from mindspore.train.amp import amp_decorator
+from mindspore.train.amp import AmpDecorator
 from mindspore._c_expression.amp import get_curr_amp_strategy
+from mindspore._check_jit_forbidden_api import jit_forbidden_register
 
 
 class _WrapCell(Cell):
@@ -104,8 +104,8 @@ class _RecomputeCell(Cell):
             set_rng_state(self.cpu_rng_state)
             _pynative_executor.set_is_run_recompute(True)
             if self.amp_strategy:
-                with amp_decorator(self.amp_strategy.get_amp_level(), self.amp_strategy.get_amp_dtype(),
-                                   self.amp_strategy.get_white_list(), self.amp_strategy.get_black_list()):
+                with AmpDecorator(self.amp_strategy.get_amp_level(), self.amp_strategy.get_amp_dtype(),
+                                  self.amp_strategy.get_white_list(), self.amp_strategy.get_black_list()):
                     grads = self.grad(self.net, self.internal_params)(*input_args, **kwargs)
             else:
                 grads = self.grad(self.net, self.internal_params)(*input_args, **kwargs)
@@ -211,14 +211,12 @@ def _detach_input(input_arg):
 def _check_validation(block):
     if not isinstance(block, Cell):
         raise TypeError("Recompute function now only support block which inherited from Cell!")
-    if context.get_context("mode") != context.PYNATIVE_MODE:
-        raise AssertionError("Recompute function now only support pynative mode, you can use "
-                             "Cell.recompute() in graph mode.")
     if block.construct.__code__.co_name == "staging_specialize":
         logger.warning('Block\'s construct method decorated by @jit that recompute '
                        'function will not come into effect.')
 
 
+@jit_forbidden_register
 def recompute(block, *args, **kwargs):
     r"""
     This function is used to reduce memory, when run block, rather than