mindspore 2.1.0__cp38-none-any.whl → 2.2.10__cp38-none-any.whl

mindspore/common/api.py

@@ -26,13 +26,14 @@ import inspect
 import importlib
 import hashlib
 import contextlib
-from collections import OrderedDict
+from collections import OrderedDict, namedtuple
 from functools import wraps
 import numpy as np
 import mindspore as ms
 from mindspore import context
 from mindspore import log as logger
 from mindspore._extends.remote import kernel_build_server
+from mindspore.common.jit_config import JitConfig
 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
@@ -48,11 +49,15 @@ from mindspore._checkparam import is_stub_tensor
 from mindspore.common._utils import is_shape_unknown
 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
 
 # Store ms_function class compiled pipeline cache.
 ms_compile_cache = set()
 # Store cell compiled pipeline cache.
 cells_compile_cache = {}
+# Store function compiled times information.
+function_phases = dict()
 
 BROADCAST_PHASE = "_broadcast_"
 _PYNATIVE_PARALLEL_FUNC_NAME = "after_shard"
@@ -79,6 +84,12 @@ def _convert_python_data(data):
     if isinstance(data, RowTensor) and not isinstance(data, PythonRowTensor):
         return PythonRowTensor(row_tensor=data)
     if isinstance(data, tuple):
+        # Handle namedtuple since its type is tuple.
+        if hasattr(data, "_fields"):
+            type_name = data.__class__.__name__
+            data_dict = data._asdict()
+            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):
         # Keep list object not change for inplace operation.
@@ -86,7 +97,11 @@ def _convert_python_data(data):
             data[i] = _convert_python_data(data[i])
         return data
     if isinstance(data, dict):
-        return dict((_convert_python_data(key), _convert_python_data(value)) for key, value in data.items())
+        # Keep the dict object not change.
+        keys = tuple(data.keys())
+        for key in keys:
+            data[_convert_python_data(key)] = _convert_python_data(data.pop(key))
+        return data
     return data
 
 
@@ -175,8 +190,7 @@ def __get_compile_cache_dep_files(file_path, compile_cache_dep_files, pkg):
         if isinstance(node, ast.ImportFrom):
             if node.module is not None:
                 module_name = node.module
-            if node.level == 1:
-                module_name = "." + module_name
+            module_name = "." * node.level + module_name
         elif not isinstance(node, ast.Import):
             continue
         # Do not care the files in mindspore package
@@ -284,195 +298,6 @@ def _get_args_for_run(obj, args, kwargs):
     return new_args
 
 
-class _AutoIdentifyDynamicShape:
-
-    """
-    Represents a function auto identify dynamic shape.
-    """
-    def __init__(self):
-        self.all_shape_cache = {}
-
-
-    @staticmethod
-    def get_input_tensor_shape(args_list):
-        """get input tensor shape and type save as tensor, and make it value to 1"""
-        tensor_list = []
-        for arg in args_list:
-            if isinstance(arg, Tensor):
-                tmp_shape = arg.shape
-                tmp_type = arg.dtype
-                tensor_list.append(PythonTensor(np.ones(tmp_shape), dtype=tmp_type))
-            else:
-                tensor_list.append(arg)
-
-        return tuple(tensor_list)
-
-
-    @staticmethod
-    def check_input_args(args_list):
-        """check input args"""
-        if not args_list:
-            return False
-
-        for elem in args_list:
-            if elem is None:
-                return False
-
-            if isinstance(elem, ms.Parameter):
-                return False
-
-            if not isinstance(elem, Tensor):
-                return False
-
-            if elem.const_arg:
-                return False
-
-        return True
-
-
-    @staticmethod
-    def _is_tensor_equal(input_tensor, cache_tensor):
-        """check two tensor is equal"""
-        if input_tensor.dtype != cache_tensor.dtype:
-            return False
-
-        if input_tensor.shape != cache_tensor.shape:
-            return False
-
-        if len(input_tensor.shape) != len(cache_tensor.shape):
-            return False
-
-        return True
-
-
-    @staticmethod
-    def _is_all_input_shape_generalize(input_shape_tuple):
-        """check all input shapes need generalize"""
-        for elem in input_shape_tuple:
-            if not is_shape_unknown(elem.shape):
-                return False
-        return True
-
-
-    def auto_dynamic_generate_compile_args(self, args_list):
-        """generate compile args in auto dynamic shape"""
-        if not self._is_enable_auto_identify_shape(args_list):
-            return args_list
-
-        args_len = len(args_list)
-        tensor_tuple = self.get_input_tensor_shape(args_list)
-        shape_cache_list = self.all_shape_cache.get(args_len)
-        # step1: init real_shape_cache, part_generalize_shape_cache, all_generalize_shape_cache.
-        if shape_cache_list is None:
-            shape_cache_list = []
-            real_shape_cache = set()
-            real_shape_cache.add(tensor_tuple)
-            shape_cache_list.append(real_shape_cache)
-            part_generalize_shape_cache = set()
-            shape_cache_list.append(part_generalize_shape_cache)
-            all_generalize_shape_cache = set()
-            shape_cache_list.append(all_generalize_shape_cache)
-            self.all_shape_cache[args_len] = shape_cache_list
-            logger.info((f'The real shape cache is empty, add it into real_shape_cache.'))
-            return tensor_tuple
-
-        # step2: find cache in real_shape_cache.
-        real_shape_cache = shape_cache_list[0]
-        is_real_shape_exist, real_shape_input = self._find_compile_args_in_shape_cache(real_shape_cache, tensor_tuple,
-                                                                                       "real")
-        if is_real_shape_exist and real_shape_input is not None:
-            return real_shape_input
-
-        # step3: if can not find cache in real_shape_cache, then generate it
-        is_generalize_shape, compile_args = self._do_generalize_shape(real_shape_cache, tensor_tuple)
-
-        # step4: if input type change or rank change, save shape into real_shape_cache and then return
-        if not is_generalize_shape and compile_args is None:
-            real_shape_cache.add(tensor_tuple)
-            return tensor_tuple
-
-        # step5: check whether all input tensor need generalize
-        all_generalize_shape_cache = shape_cache_list[2]
-        if self._is_all_input_shape_generalize(compile_args):
-            if not all_generalize_shape_cache:
-                all_generalize_shape_cache.add(compile_args)
-            logger.info((f'return all generalize shape cache.'))
-            return compile_args
-
-        # step6: find compile_args in part_generalize_shape_cache
-        part_generalize_shape_cache = shape_cache_list[1]
-        if not part_generalize_shape_cache:
-            part_generalize_shape_cache.add(compile_args)
-        else:
-            is_generalize_shape_exist, _ = self._find_compile_args_in_shape_cache(part_generalize_shape_cache,
-                                                                                  compile_args, "part generalize")
-            if not is_generalize_shape_exist:
-                logger.info((f'Can not find cache in part_generalize_shape_cache, add it into'
-                             ' part_generalize_shape_cache.'))
-                part_generalize_shape_cache.add(compile_args)
-
-        return compile_args
-
-
-    def _is_all_tensor_equal(self, input_shape_tuple, cache_shape_tuple):
-        """check two tuple is equal"""
-        for i, elem in enumerate(cache_shape_tuple):
-            res = self._is_tensor_equal(input_shape_tuple[i], elem)
-            if not res:
-                return False
-        return True
-
-
-    def _is_enable_auto_identify_shape(self, args_list):
-        """is enable auto identify shape"""
-        enable_auto_identify = os.getenv('MS_AUTO_DYNAMIC_SHAPE_ENABLE')
-        if not enable_auto_identify:
-            enable_auto_identify = False
-        if ((enable_auto_identify is False or enable_auto_identify == "0")) or not self.check_input_args(args_list):
-            return False
-        return True
-
-
-    def _find_compile_args_in_shape_cache(self, shape_cache, compile_args, cache_type):
-        """find compile args in real or part generalize shape cache"""
-        is_exist = False
-        for shapes in shape_cache:
-            is_exist = self._is_all_tensor_equal(compile_args, shapes)
-            if is_exist:
-                logger.info((f'Find cache in {cache_type} shape cache.'))
-                return is_exist, shapes
-        logger.info((f'Can not find cache in {cache_type} shape cache.'))
-        return is_exist, None
-
-
-    def _do_generalize_shape(self, real_shape_cache, tensor_tuple):
-        """do generalize shape"""
-        is_generalize_shape = False
-        for real_shape in real_shape_cache:
-            generalize_shape = []
-            for i, elem in enumerate(real_shape):
-                if len(elem.shape) != len(tensor_tuple[i].shape) or elem.dtype != tensor_tuple[i].dtype:
-                    generalize_shape.clear()
-                    break
-                if (not is_shape_unknown(elem.shape)) and self._is_tensor_equal(tensor_tuple[i], elem):
-                    generalize_shape.append(tensor_tuple[i])
-                else:
-                    shape_value = []
-                    for _ in range(len(elem.shape)):
-                        shape_value.append(-1)
-                    shape_tuple = tuple(shape_value)
-                    generalize_shape.append(PythonTensor(Tensor(shape=shape_tuple, dtype=tensor_tuple[i].dtype)))
-                    logger.info((f'The {i} input tensor shape is {tensor_tuple[i].shape}, type is '
-                                 f'{tensor_tuple[i].dtype}; in real cache shape is {elem.shape}, type is '
-                                 f'{elem.dtype}, the {i} input shape not equal, may generalize to {shape_tuple}.'))
-
-            if len(generalize_shape) == len(real_shape):
-                is_generalize_shape = True
-                return is_generalize_shape, tuple(generalize_shape)
-
-        return is_generalize_shape, None
-
-
 class _MindsporeFunctionExecutor:
     """
     Represents a function compiled by graph compiler.
@@ -490,7 +315,6 @@ 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)):
@@ -506,7 +330,7 @@ class _MindsporeFunctionExecutor:
         self._graph_executor = GraphExecutor_.get_instance()
         self._create_time = ms_create_time
         self.jit_config_dict = jit_config.jit_config_dict if jit_config else None
-        self.auto_identify_dynamic_shape = _AutoIdentifyDynamicShape()
+
 
     @_wrap_func
     def __call__(self, *args, **kwargs):
@@ -516,9 +340,9 @@ class _MindsporeFunctionExecutor:
         phase = ""
         try:
             if context.get_context("mode") == context.PYNATIVE_MODE:
-                _pynative_executor.set_ms_function_compile_status(True, phase)
+                _pynative_executor.set_jit_compile_status(True, phase)
                 phase = self.compile(self.fn.__name__, *args_list, **kwargs)
-                _pynative_executor.set_ms_function_compile_status(False, phase)
+                _pynative_executor.set_jit_compile_status(False, phase)
             else:
                 phase = self.compile(self.fn.__name__, *args_list, **kwargs)
         except Exception as err:
@@ -531,19 +355,11 @@ class _MindsporeFunctionExecutor:
         new_inputs = self._generate_run_args(args_list, kwargs)
         output = self._graph_executor(tuple(new_inputs), phase)
         if context.get_context("mode") == context.PYNATIVE_MODE:
-            output = _pynative_executor.grad_ms_function(output, *new_inputs)
-
-        enable_ge = os.getenv("MS_ENABLE_GE") == "1"
-        if enable_ge and self.jit_config_dict is None:
-            raise RuntimeError("GE and jit_level=O3 should be used together, but jit_config is None.")
-        if self.jit_config_dict:
-            enable_jit_level_o3 = self.jit_config_dict.get('jit_level') == "O3"
-            if (enable_ge and not enable_jit_level_o3) or (not enable_ge and enable_jit_level_o3):
-                raise RuntimeError("GE and jit_level=O3 should be used together, but got MS_ENABLE_GE={}, jit_level={}".
-                                   format(os.getenv("MS_ENABLE_GE"), self.jit_config_dict.get('jit_level')))
+            output = _pynative_executor.grad_jit(output, *new_inputs)
 
         return output
 
+
     def compile(self, method_name, *args, **kwargs):
         """Returns pipeline for the given args."""
         # Check whether hook function registered on Cell object.
@@ -554,14 +370,16 @@ class _MindsporeFunctionExecutor:
                                f"pynative mode and remove 'jit' decorator.")
         # Chose dynamic shape tensors or actual input tensors as compile args.
         compile_args = self._generate_compile_args(args)
+        key_id = self._get_key_id()
+        compile_args = get_auto_dynamic_shape_args_with_check_input_signature(compile_args, key_id,
+                                                                              self.input_signature)
+
         # Restore the mutable attr for every arg.
         compile_args = _restore_mutable_attr(args, compile_args)
-        generate_name = self.fn.__module__ + "." + self.fn.__name__ + "." + self.fn.__code__.co_filename + "." + \
-            str(self.fn.__code__.co_firstlineno)
-        if _pynative_executor.grad_flag():
-            generate_name = generate_name + ".grad"
-        if _is_pynative_parallel():
-            generate_name = generate_name[:generate_name.rfind(str(id(self.fn)))] + str(id(self.shard_parent_obj))
+        generate_name, echo_function_name = self._get_generate_name()
+        # The full Function name
+        full_function_name = generate_name
+        create_time = ''
 
         # Add key with obj
         if self.obj is not None:
@@ -572,13 +390,18 @@ class _MindsporeFunctionExecutor:
             self.obj.__parse_method__ = method_name
             if isinstance(self.obj, ms.nn.Cell):
                 generate_name = generate_name + '.' + str(self.obj.create_time)
+                create_time = str(self.obj.create_time)
             else:
                 generate_name = generate_name + '.' + str(self._create_time)
+                create_time = str(self._create_time)
+
             generate_name = generate_name + '.' + str(id(self.obj))
+            full_function_name = generate_name
         else:
             # Different instance of same class may use same memory(means same obj_id) at diff times.
             # To avoid unexpected phase matched, add create_time to generate_name.
             generate_name = generate_name + '.' + str(self._create_time)
+            create_time = str(self._create_time)
 
         self.enable_tuple_broaden = False
         if hasattr(self.obj, "enable_tuple_broaden"):
@@ -587,16 +410,33 @@ 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)
         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:
             return phase
 
+        self._check_recompile(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()
         if self.jit_config_dict:
             self._graph_executor.set_jit_config(self.jit_config_dict)
+        else:
+            jit_config_dict = JitConfig().jit_config_dict
+            self._graph_executor.set_jit_config(jit_config_dict)
 
         if self.obj is None:
+            # Set an attribute to fn as an identifier.
+            if isinstance(self.fn, types.MethodType):
+                setattr(self.fn.__func__, "__jit_function__", True)
+            else:
+                setattr(self.fn, "__jit_function__", True)
             is_compile = self._graph_executor.compile(self.fn, compile_args, kwargs, phase, True)
+            if isinstance(self.fn, types.MethodType):
+                delattr(self.fn.__func__, "__jit_function__")
+            else:
+                delattr(self.fn, "__jit_function__")
         else:
             if isinstance(self.obj, ms.nn.Cell):
                 self._graph_executor.set_weights_values(self.obj.parameters_dict())
@@ -605,8 +445,32 @@ class _MindsporeFunctionExecutor:
         if not is_compile:
             raise RuntimeError("Executor compile failed.")
         ms_compile_cache.add(phase)
+
         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"""
@@ -618,6 +482,31 @@ class _MindsporeFunctionExecutor:
                 opt_param.init_data()
 
 
+    def _get_key_id(self):
+        """get key id."""
+        if isinstance(self.obj, ms.nn.Cell):
+            key_id = str(id(self.obj)) + str(self.obj.create_time)
+        else:
+            key_id = str(id(self.obj)) + str(self._create_time)
+
+        if _pynative_executor.grad_flag():
+            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(
+            self.fn.__code__.co_firstlineno)
+        echo_function_name = "function \"" + self.fn.__name__ + "\" at the file \"" + self.fn.__code__.co_filename \
+                             + "\", line " + str(self.fn.__code__.co_firstlineno)
+        if _pynative_executor.grad_flag():
+            generate_name = generate_name + ".grad"
+        if _is_pynative_parallel():
+            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")
@@ -630,7 +519,7 @@ class _MindsporeFunctionExecutor:
     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 = args_list
+        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()
@@ -659,12 +548,13 @@ class _MindsporeFunctionExecutor:
                                    f"be 'sens' and added it to compile args.")
                     self.input_signature.append(args_list[-1])
                 compile_args = tuple(self.input_signature)
-                _pynative_executor.set_dynamic_input(self.obj)
+                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):
                     raise ValueError("The input args is incompatible with the args in `input_signature`!")
-        else:
-            compile_args = self.auto_identify_dynamic_shape.auto_dynamic_generate_compile_args(args_list)
         return compile_args
 
     def _generate_run_args(self, args_list, kwargs):
@@ -699,14 +589,14 @@ def _get_obj_id(input_obj):
     return obj_id + str(id(input_obj))
 
 
-def _get_ms_function_hash(hash_input):
+def _get_jit_hash(hash_input):
     """Get hash value of single object or list of objects."""
     if isinstance(list, tuple):
         return ".".join(map(_get_obj_id, hash_input))
     return _get_obj_id(hash_input)
 
 
-def jit(fn=None, input_signature=None, hash_args=None, jit_config=None):
+def jit(fn=None, input_signature=None, hash_args=None, jit_config=None, compile_once=False):
     """
     Create a callable MindSpore graph from a Python function.
 
@@ -726,6 +616,10 @@ def jit(fn=None, input_signature=None, hash_args=None, jit_config=None):
             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
+            Default: ``False`` .
 
     Returns:
         Function, if `fn` is not None, returns a callable function that will execute the compiled function; If `fn` is
@@ -769,7 +663,7 @@ def jit(fn=None, input_signature=None, hash_args=None, jit_config=None):
         ...
         >>> out = tensor_add_with_sig(x, y)
         ...
-        ... # Set hash_args as fn, otherwise cache of compiled `closure_fn` will not be reused.
+        ... # 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):
         ...     return ops.exp(x)
@@ -783,11 +677,28 @@ def jit(fn=None, input_signature=None, hash_args=None, jit_config=None):
         >>> inputs = Tensor(np.ones([10, 10, 10]).astype(np.float32))
         >>> for i in range(10):
         ...     closure_fn(inputs, func)
+        ...
+        ... # Set compile_once = True, otherwise the train_step will be compiled again.
+        >>> def train(x):
+        ...     @jit(compile_once = True)
+        ...     def train_step(x):
+        ...         return ops.exp(x)
+        ...     for i in range(10):
+        ...         train_step(x)
+        ...
+        >>> inputs = Tensor(np.ones([10, 10, 10]).astype(np.float32))
+        >>> for i in range(10):
+        ...     train(inputs)
     """
 
     def wrap_mindspore(func):
+        if not isinstance(compile_once, bool):
+            logger.warning(f"The parameter `compile_once` of jit should be a bool, "
+                           f"but got {type(compile_once)}.")
         if hash_args:
-            hash_obj = _get_ms_function_hash(hash_args)
+            hash_obj = _get_jit_hash(hash_args)
+        elif compile_once:
+            hash_obj = 0
         else:
             hash_obj = int(time.time() * 1e9)
 
@@ -984,8 +895,8 @@ def _no_recursive(callable_obj):
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
     """
-    isCellSubClass = inspect.isclass(callable_obj) and issubclass(callable_obj, ms.nn.Cell)
-    if not isCellSubClass and not inspect.ismethod(callable_obj) and not inspect.isfunction(callable_obj):
+    is_cell_subclass = inspect.isclass(callable_obj) and issubclass(callable_obj, ms.nn.Cell)
+    if not is_cell_subclass and not inspect.ismethod(callable_obj) and not inspect.isfunction(callable_obj):
         raise TypeError(f"Decorator no_recursive is used for callable object, but got {callable_obj}.")
     _add_flags(callable_obj, no_recursive=True)
     return callable_obj
@@ -1149,7 +1060,7 @@ def _build_broadcast_graph(broadcast_params_dict, broadcast_phase):
     _broadcast_net.phase = broadcast_phase
     broadcasted_params = _broadcast_net()
     for param_name, param in zip(broadcast_params_dict.keys(), broadcasted_params):
-        broadcast_params_dict[param_name].set_data(param)
+        broadcast_params_dict.get(param_name).set_data(param)
 
 
 def _get_auto_split_param_names(parameter_layout_dict):
@@ -1355,30 +1266,18 @@ class _PyNativeExecutor:
         """
         self._executor.sync()
 
-    def set_lazy_build(self, enable):
+    def grad_jit(self, output, *args):
         """
-        The switch of lazy build.
+        Building grad graph decorated by jit.
 
         Args:
-            enable (bool): Specifies whether the lazy build is enable.
+            output (tuple): The function or cell decorated by jit output object.
+            args (tuple): Function or cell decorated by jit input arguments.
 
         Return:
             None.
         """
-        self._executor.set_lazy_build(enable)
-
-    def grad_ms_function(self, output, *args):
-        """
-        Building grad graph decorated by ms_function.
-
-        Args:
-            output (tuple): The function or cell decorated by ms_function output object.
-            args (tuple): Function or cell decorated by ms_function input arguments.
-
-        Return:
-            None.
-        """
-        return self._executor.grad_ms_function(output, *args)
+        return self._executor.grad_jit(output, *args)
 
     def grad_flag(self):
         """
@@ -1422,29 +1321,42 @@ class _PyNativeExecutor:
         """
         self._executor.set_enable_grad(flag)
 
-    def set_ms_function_compile_status(self, status, phase):
+    def set_jit_compile_status(self, status, phase):
         """
-        Set ms_function is compiling
+        Set jit is compiling
 
         Args:
-            status(bool): ms_function compile status
+            status(bool): jit compile status
             phase (str): The phase of cell/function instance.
         Return:
             None.
         """
-        self._executor.set_ms_function_compile_status(status, phase)
+        self._executor.set_jit_compile_status(status, phase)
 
-    def set_dynamic_input(self, obj):
+    def set_dynamic_input(self, obj, *args):
         """
         Set dynamic shape tensor of input arguments.
 
         Args:
             obj (Function/Cell): The function or cell instance.
+            args (tuple): Function or cell dynamic input arguments.
 
         Return:
             None.
         """
-        self._executor.set_dynamic_input(obj)
+        self._executor.set_dynamic_input(obj, *args)
+
+    def get_dynamic_input(self, *actual_args):
+        """
+        Get dynamic shape arguments according to actual input arguments.
+
+        Args:
+            actual_args(tuple): Actual input arguments of Function or Cell.
+
+        Return:
+            dynamic_shape_args(tuple): Dynamic shape arguments of Function or Cell.
+        """
+        return self._executor.get_dynamic_input(*actual_args)
 
     def is_first_cell(self):
         """
@@ -1550,6 +1462,13 @@ class _CellGraphExecutor:
         """
         self._graph_executor.set_queue_name(queue_name)
 
+    def get_queue_name(self, dataset_phase):
+        """
+        Get cached queue name for the graph loaded from compile cache.
+        :return: cached queue name
+        """
+        return self._graph_executor.get_queue_name(dataset_phase)
+
     @staticmethod
     def _set_dataset_mode(obj):
         """set dataset mode."""
@@ -1597,15 +1516,18 @@ class _CellGraphExecutor:
         if not hasattr(obj, obj.__parse_method__):
             raise AttributeError(
                 'The class {} dose not have method {}'.format(obj.__class__.__name__, obj.__parse_method__))
+        key_id = str(id(obj)) + str(obj.create_time)
+        args = get_auto_dynamic_shape_args(args, key_id)
 
         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)
         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)
         phase = phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+        update_auto_dynamic_shape_phase(args, key_id, phase)
 
         if phase in obj.compile_cache and self.has_compiled(phase):
             logger.debug("%r graph has existed.", phase)
@@ -1616,12 +1538,12 @@ class _CellGraphExecutor:
         self._set_dataset_mode(obj)
         self._set_compile_cache_dep_files(phase)
 
-        enable_ge = context.get_context("enable_ge")
-        if enable_ge:
-            obj.add_flags(ge_init=True)
         self._graph_executor.set_weights_values(obj.parameters_dict())
         if jit_config_dict:
             self._graph_executor.set_jit_config(jit_config_dict)
+        else:
+            jit_config_dict = JitConfig().jit_config_dict
+            self._graph_executor.set_jit_config(jit_config_dict)
         result = self._graph_executor.compile(obj, args, kwargs, phase, self._use_vm_mode())
         obj.compile_cache.add(phase)
         if not result:
@@ -1639,17 +1561,10 @@ class _CellGraphExecutor:
             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 not do_convert:
-            return phase, True
-
-        # the following GE init process is not needed when use vm or ms backend
-        if enable_ge:
-            pass
-        elif "export" in phase:
+        if "export.air" in phase:
             self._build_data_graph(obj, phase)
         elif BROADCAST_PHASE not in phase and _get_parameter_broadcast():
             _parameter_broadcast(obj)
-
         return phase, True
 
     def _update_param_node_default_input(self, phase, replace):