mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/common/api.py

@@ -1,6 +1,6 @@
 # This is the Python adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
 #
-# Copyright 2020-2023 Huawei Technologies Co., Ltd
+# Copyright 2020-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.
@@ -38,12 +38,13 @@ from mindspore.common.tensor import Tensor as PythonTensor
 from mindspore.common.sparse_tensor import CSRTensor as PythonCSRTensor
 from mindspore.common.sparse_tensor import COOTensor as PythonCOOTensor
 from mindspore.common.sparse_tensor import RowTensor as PythonRowTensor
+from mindspore._c_expression.amp import get_curr_amp_strategy
 from mindspore._c_expression import GraphExecutor_, Tensor, CSRTensor, RowTensor, COOTensor, \
     PyNativeExecutor_, verify_inputs_signature, init_exec_dataset, _set_dataset_mode_config, init_pipeline, \
     _ms_memory_recycle, _bind_device_ctx
 from mindspore.parallel._ps_context import _is_role_sched
 from mindspore.parallel._utils import _check_full_batch, _get_parameter_broadcast, _is_pynative_parallel, \
-    _is_in_auto_parallel_mode
+    _is_in_auto_parallel_mode, _is_parallel_mode
 from mindspore import _checkparam as Validator
 from mindspore._checkparam import is_stub_tensor
 from mindspore.common._utils import is_shape_unknown
@@ -51,6 +52,8 @@ from mindspore.common.mutable import mutable
 from mindspore.common._register_for_adapter import ms_adapter_registry
 from mindspore.common.auto_dynamic_shape import get_auto_dynamic_shape_args, update_auto_dynamic_shape_phase, \
     get_auto_dynamic_shape_args_with_check_input_signature, update_auto_dynamic_shape_phase_with_check_input_signature
+from mindspore.common._pijit_context import PIJitCaptureContext
+from mindspore.common.parameter import Parameter
 
 # Store ms_function class compiled pipeline cache.
 ms_compile_cache = set()
@@ -62,6 +65,68 @@ function_phases = dict()
 BROADCAST_PHASE = "_broadcast_"
 _PYNATIVE_PARALLEL_FUNC_NAME = "after_shard"
 
+ARG_SPECIFIED = "arg_specified_infos"
+TOTAL_ARG_LEN = "total_arg_length"
+
+
+def _check_recompile_args(compile_args, kwargs):
+    """Check recompile of graph"""
+
+    def _check_constant_tensor_arg(arg):
+        if hasattr(arg, "__ms_mutable__"):
+            return False
+        if isinstance(arg, (list, tuple)):
+            return any(_check_constant_tensor_arg(x) for x in arg)
+        return isinstance(arg, Tensor)
+
+    for v in kwargs.values():
+        compile_args += (v,)
+    for arg in compile_args:
+        if not isinstance(arg, tuple) and not isinstance(arg, list):
+            continue
+        if _check_constant_tensor_arg(arg):
+            logger.warning(f"Constant value tensor are detected in tuple or list, which might cause recompiling "
+                           f"when tensor value changes. You can use mutable(Tensor) or mutable(tuple(Tensor)) "
+                           f"to set tensor's value as variable to to avoid recompiling. The tuple or list arg "
+                           f"is: {arg} .")
+            return
+
+
+def _check_recompile(obj, compile_args, kwargs, full_function_name, create_time, echo_function_name):
+    """Warning when the function has been compiled."""
+    ignore_dirs = ["mindspore/ops", "mindspore/nn"]
+    if any((lambda x: x in full_function_name)(x) for x in ignore_dirs):
+        return
+
+    if full_function_name in function_phases:
+        warning_times = 1
+        if len(function_phases[full_function_name]) >= warning_times \
+                and create_time not in function_phases[full_function_name]:
+            if isinstance(obj, ms.nn.Cell):
+                tips = f"Please try to create {echo_function_name} instance only once to avoid recompiling. "
+                logger.info(f"The {echo_function_name} has been compiled again. "
+                            f"{tips} ")
+            else:
+                tips = "Try to decorate the function with @jit(hash_args=...) " \
+                       "or @jit(compile_once=True) to reduce the compile time. " \
+                       "For more details, get instructions about `jit` at " \
+                       "https://www.mindspore.cn/search?inputValue=jit."
+                logger.warning(f"The {echo_function_name} has been compiled again. "
+                               f"{tips} ")
+        else:
+            _check_recompile_args(compile_args, kwargs)
+    else:
+        function_phases[full_function_name] = set()
+    function_phases[full_function_name].add(create_time)
+
+
+def _ms_adapter_tensor_as_parameter_output(data):
+    """Check whether the data is an output from a parameter which is a ms_adapter tensor.
+       Pylint: disable=unidiomatic-typecheck.
+    """
+    return ms_adapter_registry.is_registered and isinstance(data, ms_adapter_registry.tensor) \
+        and hasattr(data, "__ms_parameter_output__") and getattr(data, "__ms_parameter_output__")
+
 
 def _convert_python_data(data):
     """
@@ -73,8 +138,10 @@ def _convert_python_data(data):
     Returns:
         data, a data convert C++ to python
     """
-    if isinstance(data, Tensor) and data.adapter_flag:
+    if isinstance(data, (Tensor, PythonTensor)) and data.adapter_flag:
         return ms_adapter_registry.tensor(data)
+    if _ms_adapter_tensor_as_parameter_output(data) and hasattr(data, "tensor"):
+        return data.tensor
     if isinstance(data, Tensor) and not isinstance(data, PythonTensor):
         return PythonTensor(data, internal=True)
     if isinstance(data, CSRTensor) and not isinstance(data, PythonCSRTensor):
@@ -83,7 +150,7 @@ def _convert_python_data(data):
         return PythonCOOTensor(coo_tensor=data)
     if isinstance(data, RowTensor) and not isinstance(data, PythonRowTensor):
         return PythonRowTensor(row_tensor=data)
-    if isinstance(data, tuple):
+    if data.__class__ is tuple:
         # Handle namedtuple since its type is tuple.
         if hasattr(data, "_fields"):
             type_name = data.__class__.__name__
@@ -91,12 +158,12 @@ def _convert_python_data(data):
             fields = data_dict.keys()
             return namedtuple(type_name, fields)(**_convert_python_data(data_dict))
         return tuple(_convert_python_data(x) for x in data)
-    if isinstance(data, list):
+    if data.__class__ is list:
         # Keep list object not change for inplace operation.
         for i in range(len(data)):
             data[i] = _convert_python_data(data[i])
         return data
-    if isinstance(data, dict):
+    if data.__class__ is dict:
         # Keep the dict object not change.
         keys = tuple(data.keys())
         for key in keys:
@@ -167,6 +234,7 @@ def _handle_func_args(func, *args, **kwargs):
 
 sys_path = list(sys.path)
 # Get the entry script path.
+entry_script_path = None
 if sys.argv and sys.argv[0] != '':
     entry_script_path = os.path.realpath(sys.argv[0])
     entry_script_path_dir = os.path.split(entry_script_path)[0]
@@ -260,7 +328,7 @@ def _get_parameter_layout():
     return layout
 
 
-def _handle_arg(obj, arg):
+def _handle_arg(obj, arg, compile_arg):
     """Handle arg for runtime .If need handle the arg, return True"""
     if isinstance(arg, PythonTensor):
         if arg.has_init:
@@ -269,7 +337,7 @@ def _handle_arg(obj, arg):
             return arg
     elif isinstance(arg, (Tensor, CSRTensor, COOTensor)):
         return arg
-    elif hasattr(arg, "__ms_mutable__") and getattr(arg, "__ms_mutable__"):
+    elif compile_arg is not None and hasattr(compile_arg, "__ms_mutable__") and getattr(compile_arg, "__ms_mutable__"):
         # mutable([]) will be eliminated by FuncGraphSpecializer, and empty list is not supported by backend.
         if isinstance(arg, list) and not arg:
             return None
@@ -282,22 +350,185 @@ def _handle_arg(obj, arg):
     return None
 
 
-def _get_args_for_run(obj, args, kwargs):
+def _handle_arg_predict(obj, arg, compile_arg):
+    """Handle arg for runtime .If need handle the arg, return True"""
+    if arg is None:
+        return None
+
+    if isinstance(arg, (int, float)):
+        return None
+
+    if isinstance(arg, (list, tuple)):
+        if compile_arg is not None and hasattr(compile_arg, "__ms_mutable__") and \
+                getattr(compile_arg, "__ms_mutable__"):
+            # mutable([]) will be eliminated by FuncGraphSpecializer, and empty list is not supported by backend.
+            if isinstance(arg, list) and not arg:
+                return None
+            return arg
+        if hasattr(obj, "enable_tuple_broaden") and obj.enable_tuple_broaden and isinstance(arg, tuple) and \
+                _check_all_tensor(arg):
+            return arg
+        return None
+    return arg
+
+
+def _get_args_for_run(obj, args, kwargs, compile_args):
     """Get the actual input args and kwargs for runtime."""
     new_args = []
-    for arg in args:
-        new_arg = _handle_arg(obj, arg)
+    for arg, compile_arg in zip(args, compile_args):
+        new_arg = _handle_arg(obj, arg, compile_arg)
         if new_arg is not None:
             new_args.append(new_arg)
 
     for _, value in kwargs.items():
-        new_value = _handle_arg(obj, value)
+        new_value = _handle_arg(obj, value, None)
         if new_value is not None:
             new_args.append(new_value)
 
     return new_args
 
 
+def _get_args_for_run_predict(obj, args, kwargs, compile_args):
+    """Get the actual input args and kwargs for runtime."""
+    new_args = []
+    for arg, compile_arg in zip(args, compile_args):
+        new_arg = _handle_arg_predict(obj, arg, compile_arg)
+        if new_arg is not None:
+            new_args.append(new_arg)
+
+    for _, value in kwargs.items():
+        new_value = _handle_arg_predict(obj, value, None)
+        if new_value is not None:
+            new_args.append(new_value)
+
+    return new_args
+
+
+def _is_args_fullmode(args, is_init=True):
+    """Check whether the arguments is for incremental-mode.
+
+    Args:
+        args (Union[list, tuple, dict, Tensor]): Given arguments.
+        is_init (bool): Is check in argument initialization phase.
+
+    Raises:
+        RuntimeError: loss necessary keys and values for incremental-mode.
+
+    Returns:
+        bool: Fullmode or not.
+    """
+    if not isinstance(args, dict):
+        return True
+    if not is_init and (args.get(ARG_SPECIFIED, None) is None or args.get(TOTAL_ARG_LEN, None) is None):
+        raise RuntimeError(
+            "The incremental inputs should be processed(with \"%s\" and \"%s\"), but got %s." %
+            (ARG_SPECIFIED, TOTAL_ARG_LEN, str(args)))
+    return False
+
+
+def _process_dyn_args(fn, dyn_args):
+    """Process the dynamic arguments, return the necessary data for latter processing.
+
+    Args:
+        fn (Function): The root function to compile.
+        dyn_args (Union[dict, list, tuple, None]): Given arguments for dynamic compilation.
+            None for nothing, list or tuple for fullmode setting, dict for incremental configuration.
+
+    Returns:
+        A dict which contains args for dynamic compilation. None for nothing dynamic.
+    """
+    if dyn_args is None:
+        # nothing should be done for None.
+        return dyn_args
+
+    if isinstance(dyn_args, dict) and ARG_SPECIFIED in dyn_args:
+        return dyn_args
+
+    args_sig = inspect.signature(fn)
+    if _is_args_fullmode(dyn_args):
+        if not isinstance(dyn_args, (list, tuple)):
+            temp_dyn_args = (dyn_args,)
+        else:
+            temp_dyn_args = dyn_args
+
+        # If dyn_args is fullmode, it should be apply directly.
+        args_sig_parameters = list(args_sig.parameters.values())
+        if not args_sig_parameters:
+            return ()
+
+        # fn may be Cell's construct while the first input is 'self'.
+        if args_sig_parameters[0].name == "self" and (len(temp_dyn_args) + 1) == len(args_sig_parameters):
+            bound_args = args_sig.bind(None, *temp_dyn_args)
+            bound_args.apply_defaults()
+            return bound_args.args[1:]
+
+        bound_args = args_sig.bind(*temp_dyn_args)
+        bound_args.apply_defaults()
+        return bound_args.args
+
+    # The dyn_args is not fullmode, a real compilation arguments should be assembled by latter procession...
+    arg_names = []
+    args_sig_parameters = list(args_sig.parameters.values())
+    for arg_p in args_sig_parameters:
+        if arg_p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD):
+            arg_names.append(arg_p.name)
+        else:
+            raise TypeError("Dynamic arguments is not accepted for VAR_POSITIONAL or VAR_KEYWORD parameters!")
+
+    offset = -1 if fn.__name__ == 'construct' and args_sig_parameters[0].name == "self" else 0
+    meet_index = set()
+
+    def _check_index_valid(index):
+        if index >= len(arg_names):
+            raise ValueError("For dict mode, valid index is \"0\"-\"%d\", but got %s!" % (len(arg_names) - 1, index))
+        if index in meet_index:
+            raise ValueError("For dict mode, there are more than one same specified key for real index: %d!" % index)
+        meet_index.add(index)
+
+    arg_handler_infos = []
+    for k, v in dyn_args.items():
+        if not isinstance(k, str):
+            raise TypeError("For dict mode, only string key is accepted, but got %s!" % k)
+        if k in arg_names:
+            cur_id = arg_names.index(k)
+            _check_index_valid(cur_id)
+            arg_handler_infos.append([cur_id + offset, v])
+        else:
+            raise ValueError("For dict mode, valid key is %s, but got %s!" % (arg_names, k))
+    return {ARG_SPECIFIED: arg_handler_infos, TOTAL_ARG_LEN: len(args_sig_parameters)}
+
+
+def _generate_dyn_compile_args(compile_args, dyn_args):
+    """Generate the dynamic compile arguments."""
+    if not dyn_args:
+        return compile_args
+    if _is_args_fullmode(dyn_args, False):
+        if not isinstance(dyn_args, (list, tuple)):
+            return (dyn_args,)
+        return dyn_args
+    arg_specified_infos = dyn_args.get(ARG_SPECIFIED, None)
+    if arg_specified_infos is None:
+        raise RuntimeError("For dict mode, a key with \"%s\" should exist, but got %s!" %
+                           (ARG_SPECIFIED, str(dyn_args)))
+    new_compile_args = list(compile_args)
+    for index, arg in arg_specified_infos:
+        new_compile_args[index] = arg
+    return tuple(new_compile_args)
+
+
+def _get_parameter_ids(args, kwargs):
+    """Get the ids of parameters."""
+    parameter_ids = ""
+    for arg in args:
+        if isinstance(arg, Parameter):
+            parameter_ids += str(id(arg))
+    for _, value in kwargs.items():
+        # The type of key is usually String type.
+        if isinstance(value, Parameter):
+            parameter_ids += str(id(value))
+    return parameter_ids
+
+
 class _MindsporeFunctionExecutor:
     """
     Represents a function compiled by graph compiler.
@@ -315,6 +546,7 @@ class _MindsporeFunctionExecutor:
     Returns:
         The result of pipeline running in graph mode.
     """
+
     def __init__(self, fn, ms_create_time, input_signature=None, obj=None, jit_config=None):
         init_pipeline()
         if not isinstance(fn, (types.FunctionType, types.MethodType)):
@@ -329,9 +561,9 @@ class _MindsporeFunctionExecutor:
         self.enable_tuple_broaden = False
         self._graph_executor = GraphExecutor_.get_instance()
         self._create_time = ms_create_time
+        self._compile_args = None
         self.jit_config_dict = jit_config.jit_config_dict if jit_config else None
 
-
     @_wrap_func
     def __call__(self, *args, **kwargs):
         args_list = args
@@ -359,7 +591,6 @@ class _MindsporeFunctionExecutor:
 
         return output
 
-
     def compile(self, method_name, *args, **kwargs):
         """Returns pipeline for the given args."""
         # Check whether hook function registered on Cell object.
@@ -376,6 +607,7 @@ class _MindsporeFunctionExecutor:
 
         # Restore the mutable attr for every arg.
         compile_args = _restore_mutable_attr(args, compile_args)
+        self._compile_args = compile_args
         generate_name, echo_function_name = self._get_generate_name()
         # The full Function name
         full_function_name = generate_name
@@ -409,14 +641,21 @@ class _MindsporeFunctionExecutor:
 
         self._graph_executor.set_enable_tuple_broaden(self.enable_tuple_broaden)
         key = self._graph_executor.generate_arguments_key(self.fn, compile_args, kwargs, self.enable_tuple_broaden)
+
+        parameter_ids = _get_parameter_ids(args, kwargs)
+        if parameter_ids != "":
+            key = str(key) + '.' + parameter_ids
         phase = generate_name + '.' + str(key)
 
         update_auto_dynamic_shape_phase_with_check_input_signature(compile_args, key_id, phase, self.input_signature)
 
         if phase in ms_compile_cache:
+            # Release resource should be released when CompileInner won't be executed, such as cur_convert_input_
+            # generated in generate_arguments_key.
+            self._graph_executor.clear_compile_arguments_resource()
             return phase
 
-        self._check_recompile(full_function_name, create_time, echo_function_name)
+        _check_recompile(self.obj, compile_args, kwargs, full_function_name, create_time, echo_function_name)
 
         # If enable compile cache, get the dependency files list and set to graph executor.
         self._set_compile_cache_dep_files()
@@ -448,29 +687,6 @@ class _MindsporeFunctionExecutor:
 
         return phase
 
-    def _check_recompile(self, full_function_name, create_time, echo_function_name):
-        """Warning when the function has been compiled."""
-        ignore_dirs = ["mindspore/ops", "mindspore/nn"]
-        if any((lambda x: x in full_function_name)(x) for x in ignore_dirs):
-            return
-
-        if full_function_name in function_phases:
-            warning_times = 1
-            if len(function_phases[full_function_name]) >= warning_times \
-                    and create_time not in function_phases[full_function_name]:
-                tips = "Try to decorate the function with @jit(hash_args=...) " \
-                       "or @jit(compile_once=True) to reduce the compile time. " \
-                       "For more details, get instructions about `jit` at " \
-                       "https://www.mindspore.cn/search?inputValue=jit."
-
-                logger.warning(f"The {echo_function_name} has been compiled again. "
-                               f"{tips} ")
-        else:
-            function_phases[full_function_name] = set()
-
-        function_phases[full_function_name].add(create_time)
-
-
     @staticmethod
     def _optimizer_state_init(opt_states):
         """set data for all optimizer states in case it is executed in graph mode"""
@@ -481,7 +697,6 @@ class _MindsporeFunctionExecutor:
             if opt_param.has_init and (prefix in prefix_list or opt_param.name == "global_step"):
                 opt_param.init_data()
 
-
     def _get_key_id(self):
         """get key id."""
         if isinstance(self.obj, ms.nn.Cell):
@@ -493,7 +708,6 @@ class _MindsporeFunctionExecutor:
             key_id = key_id + ".grad"
         return key_id
 
-
     def _get_generate_name(self):
         """get generate name."""
         generate_name = self.fn.__module__ + "." + self.fn.__name__ + "." + self.fn.__code__.co_filename + "." + str(
@@ -506,54 +720,47 @@ class _MindsporeFunctionExecutor:
             generate_name = generate_name[:generate_name.rfind(str(id(self.fn)))] + str(id(self.shard_parent_obj))
         return generate_name, echo_function_name
 
-
     def _set_compile_cache_dep_files(self):
         # If enable compile cache, get the dependency files list
         enable_compile_cache = context.get_context("enable_compile_cache")
-        if enable_compile_cache is not True and enable_compile_cache != "1":
+        if enable_compile_cache is None:
             enable_compile_cache = os.getenv('MS_COMPILER_CACHE_ENABLE')
         if enable_compile_cache is True or enable_compile_cache == "1":
             self._graph_executor.set_compile_cache_dep_files(_get_compile_cache_dep_files())
 
-
     def _generate_compile_args(self, args_list):
         """Chose dynamic shape tensors or actual input tensors as compile args."""
         # Case: If the shape of input args is dynamic, get dynamic shape tensor from context and use it to compile.
         compile_args = _pynative_executor.get_dynamic_input(args_list)
         # Case: The `set_inputs()` of Cell object has been set, using these dynamic shape args as compile args.
         if self.fn.__name__ == 'construct' and isinstance(self.obj, ms.nn.Cell) and self.obj.get_inputs():
-            compile_args = self.obj.get_inputs()
+            compile_args = _generate_dyn_compile_args(args_list, self.obj.get_inputs())
             if len(compile_args) != len(args_list):
                 raise ValueError(f"The number of actual input tensors: {len(args_list)} is not equal to the number of "
                                  f"dynamic shape tensors: {len(compile_args)}.")
-            for i, elem in enumerate(compile_args):
-                if isinstance(elem, PythonTensor):
-                    Validator.check_dynamic_shape(compile_args[i], args_list[i], i)
+            self._graph_executor.check_argument_consistency(compile_args, args_list, "input_signature")
+            Validator.check_symbolic_shape(compile_args, args_list)
 
         # Case: If dynamic shape tensors have been assigned to `input_signature`, they are preferred as compile args.
         if self.input_signature is not None:
-            if not isinstance(self.input_signature, (tuple, list)):
-                self.input_signature = (self.input_signature,)
-            self.input_signature = list(self.input_signature)
-            dyn_shape = False
-            for i, elem in enumerate(self.input_signature):
-                if isinstance(elem, PythonTensor) and is_shape_unknown(elem.shape):
-                    Validator.check_dynamic_shape(self.input_signature[i], args_list[i], i)
-                    dyn_shape = True
+            compile_args = list(_generate_dyn_compile_args(args_list, self.input_signature))
+            dyn_shape = any([is_shape_unknown(elem.shape) for elem in compile_args if isinstance(elem, PythonTensor)])
+            Validator.check_symbolic_shape(self.input_signature, args_list)
             if dyn_shape:
                 # Checkout whether the `sens` has been added to args_list.
-                if len(self.input_signature) == len(args_list) - 1:
+                if len(compile_args) == len(args_list) - 1:
                     logger.warning(f"The number of actual input args '{len(args_list)}' is one more than the number "
-                                   f"of input_signature args '{len(self.input_signature)}'. The last actual args may "
+                                   f"of input_signature args '{len(compile_args)}'. The last actual args may "
                                    f"be 'sens' and added it to compile args.")
-                    self.input_signature.append(args_list[-1])
-                compile_args = tuple(self.input_signature)
+                    compile_args.append(args_list[-1])
+                compile_args = tuple(compile_args)
+                self._graph_executor.check_argument_consistency(compile_args, args_list, "input_signature")
                 if self.obj is not None:
                     _pynative_executor.set_dynamic_input(self.obj, *compile_args)
                 else:
                     _pynative_executor.set_dynamic_input(self.fn, *compile_args)
             else:
-                if not verify_inputs_signature(self.input_signature, args_list):
+                if not verify_inputs_signature(compile_args, args_list):
                     raise ValueError("The input args is incompatible with the args in `input_signature`!")
         return compile_args
 
@@ -568,7 +775,7 @@ class _MindsporeFunctionExecutor:
         Returns:
             new_inputs, new input args, which are required for running.
         """
-        return _get_args_for_run(self, args_list, kwargs)
+        return _get_args_for_run(self, args_list, kwargs, self._compile_args)
 
 
 # The attributes used to identify a given object.
@@ -596,29 +803,49 @@ def _get_jit_hash(hash_input):
     return _get_obj_id(hash_input)
 
 
-def jit(fn=None, input_signature=None, hash_args=None, jit_config=None, compile_once=False):
+def jit(fn=None, mode="PSJit", input_signature=None, hash_args=None, jit_config=None, compile_once=False):
     """
     Create a callable MindSpore graph from a Python function.
 
     This allows the MindSpore runtime to apply optimizations based on graph.
 
     Note:
-        If `input_signature` is specified, each input of `fn` must be a Tensor. And the input arguments for `fn`
-        will not accept `**kwargs`.
+        - If `input_signature` is specified, each input of `fn` must be a Tensor. And the input arguments for `fn`
+          will not accept `**kwargs`.
+        - It is not supported to run a function with decoration @jit(mode=“PIJit”)
+          in static graph mode, in which case the decoration @jit(mode=“PIJit”) is considered invalid.
+        - Calls to functions with decorated @jit(mode=“PIJit”) inside functions
+          decorated with @jit(mode=“PIJit”) are not supported,
+          and the decoration @jit(mode=“PIJit”) is considered invalid.
 
     Args:
         fn (Function): The Python function that will be run as a graph. Default: ``None`` .
-        input_signature (Tensor): The Tensor which describes the input arguments. The shape and dtype of the Tensor
-            will be supplied to this function. If input_signature is specified, each input to `fn` must be a `Tensor`.
-            And the input parameters of `fn` cannot accept `**kwargs`. The shape and dtype of actual inputs should
-            keep the same as input_signature. Otherwise, TypeError will be raised. Default: ``None`` .
+        mode (str): The type of jit used, the value of mode should be ``PIJit`` or ``PSJit``. Default: ``PSJit`` .
+
+            - `PSJit <https://www.mindspore.cn/docs/en/master/model_train/program_form/static_graph.html>`_ :
+              Parse python ast to build graph.
+            - `PIJit <https://www.mindspore.cn/docs/en/master/model_train/program_form/pynative.html#pijit>`_ :
+              Parse python bytecode to build graph at runtime.
+
+        input_signature (Union[Tuple, List, Dict, Tensor]): The Tensor which describes the input arguments. The
+            shape and dtype of the Tensor will be supplied to this function. If `input_signature` is specified, the
+            input parameters of `fn` cannot accept `**kwargs`, and the shape and dtype of actual inputs should keep the
+            same as `input_signature`. Otherwise, TypeError will be raised. There are two mode for `input_signature`:
+
+            - Full mode: Arguments is a Tuple, List or a Tensor, and they will be used as all compile inputs
+              for graph-compiling.
+            - Incremental mode: Argument is a Dict, and they will set to some of the graph inputs, which will be
+              substituted into the input at the corresponding position for graph-compiling.
+
+            Default: ``None`` .
+
         hash_args (Union[Object, List or Tuple of Objects]): The local free variables used inside `fn`,
             like functions or objects of class defined outside `fn`. Calling `fn` again with change of `hash_args`
             will trigger recompilation. Default: ``None`` .
         jit_config (JitConfig): Jit config for compile. Default: ``None`` .
         compile_once(bool): ``True``: The function would be compiled once when it was created many times.
             But it may be wrong if the free variables were changed. ``False`` : It would be recompiled when
-            it was created again
+            it was created again.
             Default: ``False`` .
 
     Returns:
@@ -663,6 +890,13 @@ def jit(fn=None, input_signature=None, hash_args=None, jit_config=None, compile_
         ...
         >>> out = tensor_add_with_sig(x, y)
         ...
+        >>> @jit(input_signature={"y": Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))})
+        ... def tensor_add_with_sig_1(x, y):
+        ...     z = x + y
+        ...     return z
+        ...
+        >>> out1 = tensor_add_with_sig_1(x, y)
+        ...
         ... # Set hash_args as fn, otherwise cache of compiled closure_fn will not be reused.
         ... # While fn differs during calling again, recompilation will be triggered.
         >>> def func(x):
@@ -702,6 +936,8 @@ def jit(fn=None, input_signature=None, hash_args=None, jit_config=None, compile_
         else:
             hash_obj = int(time.time() * 1e9)
 
+        dyn_args = _process_dyn_args(func, input_signature)
+
         @wraps(func)
         def staging_specialize(*args, **kwargs):
             if os.getenv("MS_JIT") == '0':
@@ -715,14 +951,24 @@ def jit(fn=None, input_signature=None, hash_args=None, jit_config=None, compile_
             # only the function or cell instance wrapped by shard will fall into this branch
             if _is_pynative_parallel() and func.__name__ == _PYNATIVE_PARALLEL_FUNC_NAME:
                 process_obj = hash_args
-            out = _MindsporeFunctionExecutor(func, hash_obj, input_signature, process_obj, jit_config)(*args, **kwargs)
+            # Handle auto mixed precision strategy.
+            if not hasattr(func, "amp_strategy"):
+                if isinstance(func, types.MethodType):
+                    setattr(func.__func__, "amp_strategy", get_curr_amp_strategy())
+                else:
+                    setattr(func, "amp_strategy", get_curr_amp_strategy())
+            out = _MindsporeFunctionExecutor(func, hash_obj, dyn_args, process_obj, jit_config)(*args, **kwargs)
             return out
 
         return staging_specialize
 
+    wrap_func = wrap_mindspore
+    if mode == "PIJit":
+        wrap_func = PIJitCaptureContext(jit_config, input_signature)
+
     if fn is not None:
-        return wrap_mindspore(fn)
-    return wrap_mindspore
+        return wrap_func(fn)
+    return wrap_func
 
 
 def ms_function(fn=None, input_signature=None, hash_args=None, jit_config=None):
@@ -732,15 +978,14 @@ def ms_function(fn=None, input_signature=None, hash_args=None, jit_config=None):
     This allows the MindSpore runtime to apply optimizations based on graph.
 
     Note:
-        `ms_function` will be deprecated and removed in a future version. Please use `jit` instead.
-        If `input_signature` is specified, each input of `fn` must be a Tensor. And the input arguments for `fn`
-        will not accept `**kwargs`.
+        - `ms_function` will be deprecated and removed in a future version. Please use :func:`mindspore.jit` instead.
+        - If `input_signature` is specified, each input of `fn` must be a Tensor. And the input arguments for `fn`
+          will not accept `**kwargs`.
 
     Args:
         fn (Function): The Python function that will be run as a graph. Default: ``None`` .
         input_signature (Tensor): The Tensor which describes the input arguments. The shape and dtype of the Tensor
-            will be supplied to this function. If input_signature is specified, each input to `fn` must be a `Tensor`.
-            And the input parameters of `fn` cannot accept `**kwargs`. The shape and dtype of actual inputs should
+            will be supplied to this function. The shape and dtype of actual inputs of `fn` should
             keep the same as input_signature. Otherwise, TypeError will be raised. Default: ``None`` .
         hash_args (Union[Object, List or Tuple of Objects]): The local free variables used inside `fn`,
             like functions or objects of class defined outside `fn`. Calling `fn` again with change of `hash_args`
@@ -909,7 +1154,7 @@ def ms_class(cls):
     This allows MindSpore to identify user-defined classes and thus obtain their attributes and methods.
 
     Note:
-        `ms_class` will be deprecated and removed in a future version. Please use `jit_class` instead.
+        `ms_class` will be deprecated and removed in a future version. Please use :func:`mindspore.jit_class` instead.
 
     Args:
         cls (Class): User-defined class.
@@ -1015,7 +1260,7 @@ def jit_class(cls):
     if not inspect.isclass(cls):
         raise TypeError(f'Decorator jit_class can only be used for class type, but got {cls}.')
     # Check if cls is nn.Cell.
-    if issubclass(cls, nn.Cell):
+    if issubclass(cls, nn.cell.Cell):
         raise TypeError(f"Decorator jit_class is used for user-defined classes and cannot be used for nn.Cell: {cls}.")
     setattr(cls, '__ms_class__', True)
     return cls
@@ -1037,6 +1282,8 @@ def set_adapter_config(config):
             ms_adapter_registry.register_parameter(value)
         elif key == "convert_object_map":
             ms_adapter_registry.register_convert_map(value)
+        elif key == "convert_adapter_tensor_map":
+            ms_adapter_registry.register_convert_adapter_tensor_map(value)
         else:
             raise ValueError(f"Unsupported key in adapter config: {key}")
 
@@ -1135,16 +1382,6 @@ class _PyNativeExecutor:
         self._executor = PyNativeExecutor_.get_instance()
         self._executor.set_py_exe_path(sys.executable)
         self._executor.set_kernel_build_server_dir(os.path.split(kernel_build_server.__file__)[0] + os.sep)
-        self._top_cell = None
-
-    def __call__(self):
-        """
-        PyNative executor run grad graph.
-
-        Return:
-            The return object after running grad graph.
-        """
-        return self._executor()
 
     @staticmethod
     def parameter_broadcast(obj, phase):
@@ -1214,23 +1451,22 @@ class _PyNativeExecutor:
         """
         self._executor.end_graph(obj, output, *args, *(kwargs.values()))
 
-    def check_run(self, grad, obj, weights, grad_hash_id, *args, **kwargs):
+    def check_run(self, grad, obj, weights, grad_hash_id, *args):
         """
         Whether the forward graph need to construct.
 
         Args:
             grad (GradOperation): The gradoperation object.
             obj (Function/Cell): The function or cell instance.
-            grad_hash_id (tuple): The id of objects which contribute to cache of compiled graph in pynative mode.
+            grad_hash_id (tuple): The id of objects, which contributes to cache of compiled graph in pynative mode.
             args (tuple): Function or cell input arguments.
-            kwargs (dict): keyword arguments.
 
         Return:
-            bool, specifies whether the forward graph need to construct.
+            bool, specifies whether the forward graph needs to construct.
         """
-        return self._executor.check_run(grad, obj, weights, grad_hash_id, *args, *(kwargs.values()))
+        return self._executor.check_run(grad, obj, weights, grad_hash_id, *args)
 
-    def grad(self, obj, grad, weights, grad_position, *args, **kwargs):
+    def grad(self, obj, grad, weights, grad_position, *args):
         """
         Get grad graph.
 
@@ -1241,12 +1477,11 @@ class _PyNativeExecutor:
             grad_position (Union(int, tuple[int])): If int, get the gradient with respect to single input.
               If tuple, get the gradients with respect to selected inputs. 'grad_position' begins with 0. Default: 0.
             args (tuple): Function or cell input arguments.
-            kwargs (dict): keyword arguments.
 
         Return:
             None.
         """
-        self._executor.grad_net(grad, obj, weights, grad_position, *args, *(kwargs.values()))
+        return self._executor.grad(grad, obj, weights, grad_position, *args)
 
     def clear_res(self):
         """
@@ -1279,9 +1514,23 @@ class _PyNativeExecutor:
         """
         return self._executor.grad_jit(output, *args)
 
+    def call_custom_bprop(self, obj, output, *args, **kwargs):
+        """
+        Call custom bprop to build variable for cell bprop.
+        Args:
+            obj (Cell): The function or cell instance.
+            output (Tensor/tuple/list): Function or cell output object.
+            args (tuple): Function or cell input arguments.
+            kwargs (dict): keyword arguments.
+
+        Return:
+            None.
+        """
+        return self._executor.call_custom_bprop(obj, output, *args, *(kwargs.values()))
+
     def grad_flag(self):
         """
-        The flag of building grad graph.
+        The flag of whether the net building grad graph.
 
         Return:
             bool, whether building grad graph.
@@ -1300,9 +1549,21 @@ class _PyNativeExecutor:
         """
         self._executor.set_grad_flag(flag)
 
+    def set_async_for_graph(self, flag):
+        """
+        Set the flag for graph async run.
+
+        Args:
+            flag (bool): Specifying whether enable graph async run.
+
+        Return:
+            None.
+        """
+        self._executor.set_async_for_graph(flag)
+
     def enable_grad(self):
         """
-        The global flag whether needing to calculate gradient.
+        The global flag that whether need to calculate gradient use in no_grad.
 
         Return:
             bool, whether needing to calculate gradient.
@@ -1321,6 +1582,18 @@ class _PyNativeExecutor:
         """
         self._executor.set_enable_grad(flag)
 
+    def requires_grad(self):
+        """
+        When both enable_grad is true and grad_flag is true, that the flag requires_grad will be true.
+
+        Args:
+            flag (bool): Specifying whether calculating gradient.
+
+        Return:
+            None.
+        """
+        return self._executor.requires_grad()
+
     def set_jit_compile_status(self, status, phase):
         """
         Set jit is compiling
@@ -1333,6 +1606,29 @@ class _PyNativeExecutor:
         """
         self._executor.set_jit_compile_status(status, phase)
 
+    def set_is_run_recompute(self, status):
+        """
+        Set recompute grad is compiling
+
+        Args:
+            status(bool): grad is in recompute status
+        Return:
+            None.
+        """
+        self._executor.set_is_run_recompute(status)
+
+    def set_cell_use_dynamic_shape_process(self, flag):
+        """
+        Set the dynamic shape flag of eval process.
+
+        Args:
+            flag (bool): Specifying whether using a dynamic process.
+
+        Return:
+            None.
+        """
+        self._executor.set_cell_use_dynamic_shape_process(flag)
+
     def set_dynamic_input(self, obj, *args):
         """
         Set dynamic shape tensor of input arguments.
@@ -1358,36 +1654,19 @@ class _PyNativeExecutor:
         """
         return self._executor.get_dynamic_input(*actual_args)
 
-    def is_first_cell(self):
-        """
-        The flag of first cell instance.
-
-        Return:
-            bool, specifies whether is the first cell.
-        """
-
-        return self._executor.is_first_cell()
-
-    def set_hook_changed(self, cell):
+    def set_mixed_precision_type(self, mixed_precision_type, is_push=True):
         """
-        The flag of registering or removing a hook function on Cell instance.
+        The value of mixed precision type.
 
         Args:
-            cell (Cell): The cell instance.
+            type(MixedPrecisionType): Mix precision type.
+            is_push(bool): If called by __enter__, is push will be True
 
         Return:
             None.
         """
-        self._executor.set_hook_changed(cell)
-
-    def get_top_cell(self):
-        """
-        Get the top cell object.
 
-        Return:
-            The top cell object.
-        """
-        return self._top_cell
+        return self._executor.set_mixed_precision_type(mixed_precision_type, is_push)
 
     def constant_folding(self, *args):
         """
@@ -1424,6 +1703,7 @@ class _CellGraphExecutor:
         self._graph_executor = GraphExecutor_.get_instance()
         self._graph_executor.set_py_exe_path(sys.executable)
         self._graph_executor.set_kernel_build_server_dir(os.path.split(kernel_build_server.__file__)[0] + os.sep)
+        self._pid = os.getpid()
 
     def init_dataset(self, queue_name, dataset_size, batch_size, dataset_types, dataset_shapes,
                      input_indexs, phase='dataset', need_run=True):
@@ -1491,9 +1771,9 @@ class _CellGraphExecutor:
     def _set_compile_cache_dep_files(self, phase):
         # If enable compile cache, get the dependency files list
         enable_compile_cache = context.get_context("enable_compile_cache")
-        if enable_compile_cache is not True and enable_compile_cache != "1":
+        if enable_compile_cache is None:
             enable_compile_cache = os.getenv('MS_COMPILER_CACHE_ENABLE')
-        if "train" in phase and (enable_compile_cache is True or enable_compile_cache == "1"):
+        if enable_compile_cache is True or enable_compile_cache == "1":
             self._graph_executor.set_compile_cache_dep_files(_get_compile_cache_dep_files())
 
     def compile(self, obj, *args, phase='predict', do_convert=True, jit_config_dict=None, **kwargs):
@@ -1522,17 +1802,30 @@ class _CellGraphExecutor:
         self.enable_tuple_broaden = False
         if hasattr(obj, "enable_tuple_broaden"):
             self.enable_tuple_broaden = obj.enable_tuple_broaden
-        logger.debug("Convert the network.", do_convert)
+        logger.debug(f"Convert the network: {do_convert}.")
         self._graph_executor.set_enable_tuple_broaden(self.enable_tuple_broaden)
         key = self._graph_executor.generate_arguments_key(obj, args, kwargs, self.enable_tuple_broaden)
         obj.arguments_key = str(key)
+        # When exist parameter in the top graph inputs, need check if the parameter object has changed.
+        parameter_ids = _get_parameter_ids(args, kwargs)
+        if parameter_ids != "":
+            obj.arguments_key = obj.arguments_key + '.' + parameter_ids
+        raw_phase = phase
         phase = phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+        obj.phase_cache[raw_phase] = phase
         update_auto_dynamic_shape_phase(args, key_id, phase)
-
+        obj.current_phase = phase
         if phase in obj.compile_cache and self.has_compiled(phase):
             logger.debug("%r graph has existed.", phase)
+            # Release resource should be released when CompileInner won't be executed, such as cur_convert_input_
+            # generated in generate_arguments_key.
+            self._graph_executor.clear_compile_arguments_resource()
             return phase, False
 
+        full_function_name = obj.__class__.__name__ + '.' + str(obj.instance_count) + '.' + str(id(type(obj)))
+        echo_function_name = obj.__class__.__name__
+        _check_recompile(obj, args, kwargs, full_function_name, obj.create_time, echo_function_name)
+
         obj.check_names()
         _check_full_batch()
         self._set_dataset_mode(obj)
@@ -1553,14 +1846,13 @@ class _CellGraphExecutor:
         if graph is None:
             raise RuntimeError("Compile graph failed for phase {}.".format(phase))
 
-        auto_parallel_mode = _is_in_auto_parallel_mode()
+        auto_parallel_mode = _is_in_auto_parallel_mode() or _is_parallel_mode()
         if not auto_parallel_mode:
             replace = obj.init_parameters_data(auto_parallel_mode=auto_parallel_mode)
             self._update_param_node_default_input(phase, replace)
         elif 'skip_auto_parallel_compile' not in obj.get_flags().keys():
             obj.parameter_layout_dict = self._graph_executor.get_parameter_layout(phase)
             obj.parallel_parameter_name_list = self._graph_executor.get_parallel_parameter_name_list(phase)
-
         if "export.air" in phase:
             self._build_data_graph(obj, phase)
         elif BROADCAST_PHASE not in phase and _get_parameter_broadcast():
@@ -1600,6 +1892,18 @@ class _CellGraphExecutor:
         """
         return self._graph_executor.has_compiled(phase)
 
+    def flops_collection(self, phase='train'):
+        """
+        Specify whether have been compiled.
+
+        Args:
+            phase (str): The phase name. Default: 'predict'.
+
+        Returns:
+            bool, specifies whether the specific graph has been compiled.
+        """
+        return self._graph_executor.flops_collection(phase)
+
     @_wrap_func
     def _exec_pip(self, obj, *args, phase=''):
         """Execute the generated pipeline."""
@@ -1630,7 +1934,9 @@ class _CellGraphExecutor:
 
     def del_net_res(self, obj, net_id):
         """Clear the memory resource of a network."""
-        self._graph_executor.del_net_res(obj, net_id)
+        # no need to del net res by gc in independent dataset process which is a subprocess forked by main process
+        if self._pid == os.getpid():
+            self._graph_executor.del_net_res(obj, net_id)
 
     def _get_branch_control_input(self):
         if ('obf_ratio' not in self.obfuscate_config.keys()) or (
@@ -1738,7 +2044,21 @@ def _bind_device_context():
     _bind_device_ctx()
 
 
+def flops_collection(phase='train'):
+    """
+    Recycle memory used by MindSpore.
+    When train multi Neural network models in one process, memory used by MindSpore is very large,
+    this is because MindSpore cached runtime memory for every model.
+    To recycle these cached memory, users can call this function after training of one model.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.ms_memory_recycle()
+    """
+    return _cell_graph_executor.flops_collection(phase)
+
+
 _cell_graph_executor = _CellGraphExecutor()
 _pynative_executor = _PyNativeExecutor()
 
-__all__ = ['ms_function', 'ms_memory_recycle', 'ms_class', 'jit', 'jit_class']
+__all__ = ['ms_function', 'ms_memory_recycle', 'ms_class', 'jit', 'jit_class', 'flops_collection']