mindspore 2.1.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.0__cp38-cp38-manylinux1_x86_64.whl

mindspore/nn/cell.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2022 Huawei Technologies Co., Ltd
+# Copyright 2020-2023 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.
@@ -33,8 +33,7 @@ from mindspore import context
 from mindspore._c_expression import init_pipeline, update_func_graph_hyper_params, Cell_, FuncGraph, MixedPrecisionType
 from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
-from mindspore.common.api import _cell_graph_executor, _pynative_executor, _get_args_for_run, cells_compile_cache, \
-    _AutoIdentifyDynamicShape
+from mindspore.common.api import _cell_graph_executor, _pynative_executor, _get_args_for_run, cells_compile_cache
 from mindspore.common.api import _generate_branch_control_input
 from mindspore.common.parameter import Parameter, ParameterTuple
 from mindspore.common.tensor import Tensor
@@ -65,6 +64,15 @@ class Cell(Cell_):
     graph in GRAPH_MODE (static graph mode) and used as the basic module of neural networks in
     PYNATIVE_MODE (dynamic graph mode).
 
+    .. note::
+        Cell is the inference mode by default. For a class that inherits a Cell,
+        if the training and inference have different structures, the subclass performs the inference branch by default.
+        To set the training mode, refer to `mindspore.nn.Cell.set_train` .
+
+    .. warning::
+        In the subclass of Cell, it's not allowed to define a method named 'cast' and not allowed to define an attribute
+        named 'phase' or 'cells', otherwise, an error will be raised.
+
     Args:
         auto_prefix (bool, optional): Whether to automatically generate NameSpace for Cell and its child cells. It also
                       affects the names of parameters in the `Cell`. If set to ``True`` , the parameter name will be
@@ -156,11 +164,9 @@ class Cell(Cell_):
         self.saved_dynamic_shape = None
         self._jit_config_dict = dict()
         self.grad_ops_label = False
-        self.to_float_fp16 = False
-        self.ge_init = False
         self.ge_sync_data = False
-        self.auto_identify_dynamic_shape = _AutoIdentifyDynamicShape()
-
+        self._is_check_and_refresh = False
+        self._amp_level = ""
 
     def __getstate__(self):
         base = Cell_.__getstate__(self)
@@ -192,6 +198,23 @@ class Cell(Cell_):
     def param_prefix(self):
         """
         Param prefix is the prefix of current cell's direct child parameter.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.dense = nn.Dense(2, 2)
+            ...
+            ...     def construct(self, x):
+            ...         x = self.dense(x)
+            ...         return x
+            >>> net = Net()
+            >>> net.update_cell_prefix()
+            >>> print(net.dense.param_prefix)
+            dense
         """
         return self._param_prefix
 
@@ -202,7 +225,7 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Cell and Parameter - Custom Cell Reverse
-              <https://mindspore.cn/tutorials/en/r2.1/advanced/modules/layer.html#custom-cell-reverse>`_
+              <https://mindspore.cn/tutorials/en/r2.2/advanced/modules/layer.html#custom-cell-reverse>`_
         """
         return self._bprop_debug
 
@@ -309,6 +332,21 @@ class Cell(Cell_):
         for item in self.trainable_params():
             item.add_pipeline_stage(value)
 
+    @property
+    def pipeline_segment(self):
+        return self._pipeline_segment
+
+    @pipeline_segment.setter
+    def pipeline_segment(self, value):
+        if not isinstance(value, int) or isinstance(value, bool):
+            raise TypeError("For 'context.set_auto_parallel_context', the argument 'pipeline_stages' "
+                            "must be int type, but got type : {}".format(type(value)))
+
+        if value < 0:
+            raise ValueError("For 'context.set_auto_parallel_context', the argument 'pipeline_stages' "
+                             "can not be less than 0, but got {}".format(value))
+        self._pipeline_segment = value
+
     @property
     def parallel_parameter_merge_net_dict(self):
         return self._parallel_parameter_merge_net_dict
@@ -345,7 +383,7 @@ class Cell(Cell_):
         if '_params_list' in self.__dict__:
             params_list = self.__dict__['_params_list']
             if name in params_list:
-                return ParameterTuple(params_list[name])
+                return params_list[name]
         raise AttributeError("The '{}' object has no attribute '{}'.".format(type(self).__name__, name))
 
     def __del__(self):
@@ -365,11 +403,11 @@ class Cell(Cell_):
             del self._params[name]
         elif name in self._cells:
             del self._cells[name]
+        elif '_params_list' in self.__dict__ and name in self._params_list:
+            del self._params_list[name]
+        elif '_tensor_list' in self.__dict__ and name in self._tensor_list:
+            del self._tensor_list[name]
         else:
-            if '_params_list' in self.__dict__ and name in self._params_list:
-                del self._params_list[name]
-            elif '_tensor_list' in self.__dict__ and name in self._tensor_list:
-                del self._tensor_list[name]
             object.__delattr__(self, name)
         self._attr_synced = False
 
@@ -381,8 +419,8 @@ class Cell(Cell_):
                 res.append(self._cast_mixed_precision_inputs(item, dst_type))
             elif isinstance(item, float):
                 res.append(self.cast(item, dst_type))
-            elif hasattr(item, "dtype") and item.dtype in {mstype.float16, mstype.float32, mstype.float64} and \
-                item.dtype != dst_type:
+            elif hasattr(item, "dtype") and item.dtype in \
+                {mstype.float16, mstype.float32, mstype.float64, mstype.bfloat16} and item.dtype != dst_type:
                 res.append(self.cast(item, dst_type))
             else:
                 res.append(item)
@@ -629,7 +667,10 @@ class Cell(Cell_):
         if PackFunc.is_tracing():
             return self._run_tracefunc(*args, **kwargs)
 
-        self.check_names_and_refresh_name()
+        if hasattr(self, '_is_check_and_refresh') and not self._is_check_and_refresh:
+            self.check_names_and_refresh_name()
+            self._is_check_and_refresh = True
+
         # Run in Graph mode.
         if os.getenv("MS_JIT") != '0' and context._get_mode() == context.GRAPH_MODE:
             self._check_construct_args(*args)
@@ -886,14 +927,14 @@ class Cell(Cell_):
             >>> import mindspore as ms
             >>> from mindspore import nn, Tensor
             >>>
-            >>> class reluNet(nn.Cell):
+            >>> class ReluNet(nn.Cell):
             ...     def __init__(self):
-            ...         super(reluNet, self).__init__()
+            ...         super(ReluNet, self).__init__()
             ...         self.relu = nn.ReLU()
             ...     def construct(self, x):
             ...         return self.relu(x)
             >>>
-            >>> net = reluNet()
+            >>> net = ReluNet()
             >>> input_dyn = Tensor(shape=[3, None], dtype=ms.float32)
             >>> net.set_inputs(input_dyn)
             >>> input1 = Tensor(np.random.random([3, 10]), dtype=ms.float32)
@@ -902,13 +943,10 @@ class Cell(Cell_):
         if self.grad_ops_label:
             logger.warning(f'For Cell, set_inputs must be set before the gradient function of the network is '
                            f'generated.')
-        for ele in inputs:
-            if isinstance(ele, str):
-                raise TypeError(f"For element in 'set_inputs', the type must not be str.")
         self._dynamic_shape_inputs = inputs
         self._check_construct_args(*inputs)
         if context._get_mode() == context.PYNATIVE_MODE:
-            _pynative_executor.set_dynamic_input(self)
+            _pynative_executor.set_dynamic_input(self, *self._dynamic_shape_inputs)
 
     def get_inputs(self):
         """
@@ -919,6 +957,26 @@ class Cell(Cell_):
 
         .. warning::
             This is an experimental API that is subject to change or deletion.
+
+        Examples:
+            >>> import numpy as np
+            >>> import mindspore as ms
+            >>> from mindspore import nn, Tensor
+            >>>
+            >>> class ReluNet(nn.Cell):
+            ...     def __init__(self):
+            ...         super(ReluNet, self).__init__()
+            ...         self.relu = nn.ReLU()
+            ...     def construct(self, x):
+            ...         return self.relu(x)
+            >>>
+            >>> net = ReluNet()
+            >>> input_dyn = Tensor(shape=[3, None], dtype=ms.float32)
+            >>> net.set_inputs(input_dyn)
+            >>> get_inputs = net.get_inputs()
+            >>> print(get_inputs)
+            (Tensor(shape=[3, -1], dtype=Float32, value= ),)
+
         """
 
         return self._dynamic_shape_inputs
@@ -936,9 +994,8 @@ class Cell(Cell_):
             self._dynamic_shape_inputs = convert_inputs_to_dynamic(*args)
 
         if self._dynamic_shape_inputs is None:
-            compile_args = self.auto_identify_dynamic_shape.auto_dynamic_generate_compile_args(args)
             _cell_graph_executor.compile(self, phase=self.phase,
-                                         jit_config_dict=self._jit_config_dict, *compile_args, **kwargs)
+                                         jit_config_dict=self._jit_config_dict, *args, **kwargs)
         else:
             self._check_compile_dynamic_shape(self._dynamic_shape_inputs, args)
             self.saved_dynamic_shape = self._dynamic_shape_inputs
@@ -994,6 +1051,23 @@ class Cell(Cell_):
         Raises:
             KeyError: If the name of parameter is null or contains dot.
             TypeError: If the type of parameter is not Parameter.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn, Parameter
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.relu = nn.ReLU()
+            ...
+            ...     def construct(self, x):
+            ...         x = self.relu(x)
+            ...         return x
+            >>> net = Net()
+            >>> net.insert_param_to_cell("bias", Parameter(Tensor([1, 2, 3])))
+            >>> print(net.bias)
+            Parameter(name=bias, shape=(3,), dtype=Int64, requires_grad=True)
         """
         if not param_name:
             raise KeyError("For 'insert_param_to_cell', the argument 'param_name' should not be None.")
@@ -1007,6 +1081,9 @@ class Cell(Cell_):
         if not isinstance(param, Parameter) and param is not None:
             raise TypeError(f"For 'insert_param_to_cell', the argument 'param' must be 'Parameter' if not None, "
                             f"but got {type(param)}.")
+        if param is None:
+            raise TypeError(f"For 'insert_param_to_cell', the argument 'param' must not be None, "
+                            f"but got None.")
         if isinstance(param, Parameter) and param.name == PARAMETER_NAME_DEFAULT:
             param.name = param_name
         self._params[param_name] = param
@@ -1048,6 +1125,18 @@ class Cell(Cell_):
             KeyError: Child Cell's name is incorrect or duplicated with the other child name.
             TypeError: If type of `child_name` is not str.
             TypeError: Child Cell's type is incorrect.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn
+            ...
+            >>> net1 = nn.ReLU()
+            >>> net2 = nn.Dense(2, 2)
+            >>> net1.insert_child_to_cell("child", net2)
+            >>> print(net1)
+            ReLU<
+              (child): Dense<input_channels=2, output_channels=2, has_bias=True>
+              >
         """
         if not isinstance(child_name, str):
             raise TypeError(f"For 'insert_child_to_cell', the type of parameter 'child_name' must be str, "
@@ -1118,6 +1207,25 @@ class Cell(Cell_):
 
         Returns:
             Dict[Parameter, Parameter], returns a dict of original parameter and replaced parameter.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.dense = nn.Dense(2, 2)
+            ...
+            ...     def construct(self, x):
+            ...         x = self.dense(x)
+            ...         return x
+            >>> net = Net()
+            >>> print(net.init_parameters_data())
+            {Parameter (name=dense.weight, shape=(2,2), dtype=Float32, requires_grad=True):
+             Parameter (name=dense.weight, shape=(2,2), dtype=Float32, requires_grad=True),
+             Parameter (name=dense.bias, shape=(2,), dtype=Float32, requires_grad=True):
+             Parameter (name=dense.bias, shape=(2,), dtype=Float32, requires_grad=True)}
         """
         replace = dict()
 
@@ -1163,6 +1271,24 @@ class Cell(Cell_):
 
         Returns:
             OrderedDict, return parameters dictionary.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn, Parameter
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.dense = nn.Dense(2, 2)
+            ...
+            ...     def construct(self, x):
+            ...         x = self.dense(x)
+            ...         return x
+            >>> net = Net()
+            >>> print(net.parameters_dict())
+            OrderedDict([('dense.weight', Parameter(name=dense.weight, shape=(2, 2), dtype=Float32,
+            requires_grad=True)), ('dense.bias', Parameter(name=dense.bias, shape=(2,), dtype=Float32,
+            requires_grad=True))])
         """
         param_dict = OrderedDict()
         for param in self.get_parameters(expand=recurse):
@@ -1238,7 +1364,7 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Model Training - Optimizer
-              <https://mindspore.cn/tutorials/en/r2.1/beginner/train.html#optimizer>`_
+              <https://mindspore.cn/tutorials/en/r2.2/beginner/train.html#optimizer>`_
         """
         return list(filter(lambda x: x.requires_grad, self.get_parameters(expand=recurse)))
 
@@ -1263,6 +1389,7 @@ class Cell(Cell_):
         Returns an iterator over cell parameters.
 
         Yields parameters of this cell. If `expand` is ``true`` , yield parameters of this cell and all subcells.
+        For more details about subcells, please see the example below.
 
         Args:
             expand (bool): If ``true`` , yields parameters of this cell and all subcells. Otherwise, only yield
@@ -1272,11 +1399,34 @@ class Cell(Cell_):
             Iteration, all parameters at the cell.
 
         Examples:
-            >>> from mindspore import nn
-            >>> net = nn.Dense(3, 4)
-            >>> parameters = []
-            >>> for item in net.get_parameters():
-            ...     parameters.append(item)
+            >>> import mindspore as ms
+            >>> from mindspore import nn, ops, Tensor
+            >>> import numpy as np
+            >>> class TestNet(nn.Cell):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.my_w1 = ms.Parameter(Tensor(np.ones([4, 4]), ms.float32))
+            ...         self.my_w2 = ms.Parameter(Tensor(np.ones([16]), ms.float32))
+            ...     def construct(self, x):
+            ...         x += self.my_w1
+            ...         x = ops.reshape(x, (16,)) - self.my_w2
+            ...         return x
+            >>> class TestNet2(nn.Cell):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.my_t1 = ms.Parameter(Tensor(np.ones([4, 4]), ms.float32))
+            ...         # self.subcell is a subcell of TestNet2, when using expand=True, the parameters of TestNet will
+            ...         # also be gathered.
+            ...         self.subcell = TestNet()
+            ...     def construct(self, x):
+            ...         x += self.my_w1
+            ...         x = ops.reshape(x, (16,)) - self.my_w2
+            ...         return x
+            >>> net = TestNet2()
+            >>> print([p for p in net.get_parameters(expand=True)])
+            [Parameter (name=my_t1, shape=(4, 4), dtype=Float32, requires_grad=True), Parameter (name=subcell.my_w1,
+            shape=(4, 4), dtype=Float32, requires_grad=True), Parameter (name=subcell.my_w2, shape=(16,), dtype=Float32,
+            requires_grad=True)]
         """
         for _, param in self.parameters_and_names(expand=expand):
             yield param
@@ -1325,7 +1475,7 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Building a Network - Model Parameters
-              <https://mindspore.cn/tutorials/en/r2.1/beginner/model.html#model-parameters>`_
+              <https://mindspore.cn/tutorials/en/r2.2/beginner/model.html#model-parameters>`_
         """
         cells = []
         if expand:
@@ -1337,7 +1487,7 @@ class Cell(Cell_):
         for cell_name, cell in cells:
             params = cell._params.items()
             for par_name, par in params:
-                if par.inited_param is not None:
+                if par is not None and par.inited_param is not None:
                     par = par.inited_param
                 if par is not None and id(par) not in params_set:
                     params_set.add(id(par))
@@ -1394,6 +1544,22 @@ class Cell(Cell_):
 
         Returns:
             Iteration, the immediate cells in the cell.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.dense = nn.Dense(2, 2)
+            ...
+            ...     def construct(self, x):
+            ...         x = self.dense(x)
+            ...         return x
+            >>> net = Net()
+            >>> print(net.cells())
+            odict_values([Dense<input_channels=2, output_channels=2, has_bias=True>])
         """
         return self.name_cells().values()
 
@@ -1439,6 +1605,22 @@ class Cell(Cell_):
 
         Returns:
             Dict, all the child cells and corresponding names in the cell.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.dense = nn.Dense(2, 2)
+            ...
+            ...     def construct(self, x):
+            ...         x = self.dense(x)
+            ...         return x
+            >>> net = Net()
+            >>> print(net.name_cells())
+            OrderedDict([('dense', Dense<input_channels=2, output_channels=2, has_bias=True>)])
         """
         value_set = set()
         cells = OrderedDict()
@@ -1454,13 +1636,8 @@ class Cell(Cell_):
             Cell_.set_mixed_precision_type(self, MixedPrecisionType.FP16)
         if "fp32" in flags and flags.get("fp32", False):
             Cell_.set_mixed_precision_type(self, MixedPrecisionType.FP32)
-
-    def _add_mixed_precision_flag_recursive(self, **flags):
-        """Add mixed precision flag to each cell"""
-        if "fp16" in flags and flags.get("fp16", False):
-            self._set_mixed_precision_type_recursive(MixedPrecisionType.FP16)
-        if "fp32" in flags and flags.get("fp32", False):
-            self._set_mixed_precision_type_recursive(MixedPrecisionType.FP32)
+        if "bf16" in flags and flags.get("bf16", False):
+            Cell_.set_mixed_precision_type(self, MixedPrecisionType.BF16)
 
     def apply(self, fn):
         """
@@ -1503,6 +1680,23 @@ class Cell(Cell_):
         Args:
             flags (dict): Network configuration information, currently it is used for the binding of network and
                 dataset. Users can also customize network attributes by this parameter.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.relu = nn.ReLU()
+            ...
+            ...     def construct(self, x):
+            ...         x = self.relu(x)
+            ...         return x
+            >>> net = Net()
+            >>> net.add_flags(sink_mode=True)
+            >>> print(net.sink_mode)
+            True
         """
         if not hasattr(self, "_func_graph_flags"):
             self._func_graph_flags = {}
@@ -1518,9 +1712,25 @@ class Cell(Cell_):
         Args:
             flags (dict): Network configuration information, currently it is used for the binding of network and
                 dataset. Users can also customize network attributes by this parameter.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.relu = nn.ReLU()
+            ...
+            ...     def construct(self, x):
+            ...         x = self.relu(x)
+            ...         return x
+            >>> net = Net()
+            >>> net.add_flags_recursive(sink_mode=True)
+            >>> print(net.sink_mode)
+            True
         """
         self.add_flags(**flags)
-        self._add_mixed_precision_flag_recursive(**flags)
         for cell in self.cells():
             cell.add_flags_recursive(**flags)
         return self
@@ -1532,17 +1742,28 @@ class Cell(Cell_):
     def get_flags(self):
         """
         Get the self_defined attributes of the cell, which can be added by `add_flags` method.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.relu = nn.ReLU()
+            ...
+            ...     def construct(self, x):
+            ...         x = self.relu(x)
+            ...         return x
+            >>> net = Net()
+            >>> net.add_flags(sink_mode=True)
+            >>> print(net.get_flags())
+            {'sink_mode':True}
         """
         if not hasattr(self, "_func_graph_flags"):
             self._func_graph_flags = {}
         return self._func_graph_flags
 
-    def _set_mixed_precision_type_recursive(self, mixed_type):
-        """Set mixed precision type to each cell"""
-        Cell_.set_mixed_precision_type(self, mixed_type)
-        for cell in self.cells():
-            cell._set_mixed_precision_type_recursive(mixed_type)
-
     def to_float(self, dst_type):
         """
         Add cast on all inputs of cell and child cells to run with certain float type.
@@ -1555,13 +1776,13 @@ class Cell(Cell_):
 
         Args:
             dst_type (:class:`mindspore.dtype`): Transfer cell to run with dst_type.
-                dst_type can be `mstype.float16` or `mstype.float32`.
+                dst_type can be `mstype.float16` , `mstype.float32` or `mstype.bfloat16`.
 
         Returns:
             Cell, the cell itself.
 
         Raises:
-            ValueError: If dst_type is not mstype.float32 or mstype.float16.
+            ValueError: If dst_type is not `mstype.float32` , `mstype.float16` or `mstype.bfloat16`.
 
         Supported Platforms:
             ``Ascend`` ``GPU`` ``CPU``
@@ -1573,19 +1794,15 @@ class Cell(Cell_):
             >>> net = nn.Conv2d(120, 240, 4, has_bias=False, weight_init='normal')
             >>> net.to_float(mstype.float16)
             Conv2d<input_channels=120, output_channels=240, kernel_size=(4, 4), stride=(1, 1), pad_mode=same,
-            padding=0, dilation=(1, 1), group=1, has_bias=False, weight_init=normal, bias_init=zeros, format=NCHW>
-        """
-        if dst_type not in (mstype.float16, mstype.float32):
-            raise ValueError("For 'to_float', the argument 'dst_type' must be mstype.float32 or mstype.float16, "
-                             "but got type: {} and value: {}.".format(type(dst_type), dst_type))
-        if dst_type == mstype.float16:
-            self._set_mixed_precision_type_recursive(MixedPrecisionType.FP16)
-            self.to_float_fp16 = True
-        else:
-            self._set_mixed_precision_type_recursive(MixedPrecisionType.FP32)
-            self.to_float_fp16 = False
-        flags = {'fp16': dst_type == mstype.float16, 'fp32': dst_type == mstype.float32}
+            padding=0, dilation=(1, 1), group=1, has_bias=False, weight_init=normal, bias_init=None, format=NCHW>
+        """
+        if dst_type not in (mstype.float16, mstype.float32, mstype.bfloat16):
+            raise ValueError("For 'to_float', the argument 'dst_type' must be mstype.float32, mstype.float16 or "
+                             "mstype.bfloat16, but got type: {} and value: {}.".format(type(dst_type), dst_type))
+        flags = {'fp16': dst_type == mstype.float16, 'fp32': dst_type == mstype.float32,
+                 'bf16': dst_type == mstype.bfloat16}
         self._add_init_args(**flags)
+        self.add_flags_recursive(**flags)
         return self
 
     def set_boost(self, boost_type):
@@ -1594,7 +1811,7 @@ class Cell(Cell_):
         accelerate the algorithm in the algorithm library.
 
         If `boost_type` is not in the algorithm library, please view the algorithm in the algorithm library through
-        `algorithm library <https://gitee.com/mindspore/mindspore/tree/r2.1/mindspore/python/mindspore/boost>`_.
+        `algorithm library <https://gitee.com/mindspore/mindspore/tree/r2.2/mindspore/python/mindspore/boost>`_.
 
         Note:
             Some acceleration algorithms may affect the accuracy of the network, please choose carefully.
@@ -1651,12 +1868,12 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Model Training - Implementing Training and Evaluation
-              <https://mindspore.cn/tutorials/en/r2.1/beginner/train.html#implementing-training-and-evaluation>`_
+              <https://mindspore.cn/tutorials/en/r2.2/beginner/train.html#training-and-evaluation>`_
         """
-        if mode is False:
-            self._phase = 'predict'
-        else:
+        if mode:
             self._phase = 'train'
+        else:
+            self._phase = 'predict'
         self.add_flags_recursive(training=mode)
         return self
 
@@ -1685,16 +1902,27 @@ class Cell(Cell_):
 
         Args:
             jit_config (JitConfig): Jit config for compile. For details, please refer to :class:`mindspore.JitConfig`.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.relu = nn.ReLU()
+            ...
+            ...     def construct(self, x):
+            ...         x = self.relu(x)
+            ...         return x
+            >>> net = Net()
+            >>> jitconfig = ms.JitConfig()
+            >>> net.set_jit_config(jitconfig)
         """
         if self._jit_config_dict:
             logger.warning("For Cell, jit config can only be set once, ignore this setting.")
         else:
             self._jit_config_dict = jit_config.jit_config_dict
-            enable_ge = os.getenv("MS_ENABLE_GE") == '1'
-            enable_jit_level_o3 = self._jit_config_dict.get('jit_level') == "O3"
-            if (not enable_ge and enable_jit_level_o3) or (enable_ge and not enable_jit_level_o3):
-                raise RuntimeError("GE and jit_level=O3 should be used together, but got MS_ENABLE_GE={}, jie_level={}".
-                                   format(os.getenv("MS_ENABLE_GE"), self.jit_config_dict.get('jit_level')))
 
     def flatten_weights(self, fusion_size=0):
         """
@@ -2290,12 +2518,13 @@ class Cell(Cell_):
     def _run_tracefunc(self, *args, **kwargs):
         """ Run Packed Cell in Pack."""
         args = self._mixed_precision_cast(args)
-        if hasattr(self, "bprop") or hasattr(self, "_pipeline_stage") or self.get_flags():
+        need_subgraph = hasattr(self, "bprop") or hasattr(self, "_pipeline_stage") or self.get_flags()
+        if not PackFunc.current.is_pynative_mode and need_subgraph:
             expander = PackExpander.get_instance()
             args = expander.begin_subgraph(self, *args)
             args = [_convert_tensor(a) for a in args]
             output = self._run_construct(args, kwargs)
-            ret = expander.end_subgraph(output)
+            ret = expander.end_subgraph(self, output)
             output = _convert_tensor(ret)
         else:
             with _SetMixedPrecision(self):
@@ -2306,10 +2535,23 @@ class Cell(Cell_):
         mixed_type = self.get_mixed_precision_type()
         if mixed_type == MixedPrecisionType.NOTSET:
             return inputs
-        cast_type = mstype.float16 if mixed_type == MixedPrecisionType.FP16 else mstype.float32
+        if mixed_type == MixedPrecisionType.FP16:
+            cast_type = mstype.float16
+        elif mixed_type == MixedPrecisionType.BF16:
+            cast_type = mstype.bfloat16
+        else:
+            cast_type = mstype.float32
         cast_inputs = self._cast_mixed_precision_inputs(inputs, cast_type)
         return cast_inputs
 
+    def _get_attr_from_cell(self, network):
+        if not isinstance(network, Cell):
+            return
+        if hasattr(network, "jit_config_dict"):
+            self._jit_config_dict = network.jit_config_dict
+        if hasattr(network, "_amp_level"):
+            self._amp_level = getattr(network, "_amp_level")
+
 
 class GraphCell(Cell):
     """