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

mindspore/nn/probability/distribution/poisson.py

@@ -36,7 +36,8 @@ class Poisson(Distribution):
     where :math:`\lambda` is the rate of the distribution.
 
     Args:
-        rate (list, numpy.ndarray, Tensor): The rate of the Poisson distribution. :math:`\lambda` in the formula. Default: ``None`` .
+        rate (list, numpy.ndarray, Tensor): The rate of the Poisson distribution.
+            :math:`\lambda` in the formula. Default: ``None`` .
         seed (int): The seed used in sampling. The global seed is used if it is ``None`` . Default: ``None`` .
         dtype (mindspore.dtype): The type of the event samples. Default: ``mstype.float32`` .
         name (str): The name of the distribution. Default: ``'Poisson'`` .

mindspore/nn/sparse/sparse.py

@@ -15,8 +15,8 @@
 """Sparse related tools."""
 from __future__ import absolute_import
 
+from mindspore import ops
 from mindspore import log as logger
-from mindspore.ops import operations as P
 from mindspore.nn.cell import Cell
 
 
@@ -76,7 +76,7 @@ class SparseToDense(Cell):
         logger.warning("'nn.SparseToDense' is deprecated from version 2.0 and will be removed in a future version. " +
                        "Please use 'COOTensor.to_dense()' instead.")
         super(SparseToDense, self).__init__()
-        self.sparse_to_dense = P.SparseToDense()
+        self.sparse_to_dense = ops.SparseToDense()
 
     def construct(self, sparse_tensor):
         return self.sparse_to_dense(sparse_tensor.indices,
@@ -141,7 +141,7 @@ class SparseTensorDenseMatmul(Cell):
         super(SparseTensorDenseMatmul, self).__init__()
         self.adj_st = adjoint_st
         self.adj_dt = adjoint_dt
-        self.sparse_dense_matmul = P.SparseTensorDenseMatmul(adjoint_st=self.adj_st, adjoint_dt=self.adj_dt)
+        self.sparse_dense_matmul = ops.SparseTensorDenseMatmul(adjoint_st=self.adj_st, adjoint_dt=self.adj_dt)
 
     def construct(self, indices, values, sparse_shape, dense):
         return self.sparse_dense_matmul(indices, values, sparse_shape, dense)

mindspore/nn/wrap/cell_wrapper.py

@@ -23,7 +23,7 @@ from types import FunctionType, MethodType
 
 from mindspore import log as logger
 from mindspore.parallel._utils import _get_device_num, _get_gradients_mean,\
-    _get_parallel_mode, _get_enable_parallel_optimizer, _is_pynative_parallel
+    _get_parallel_mode, _get_enable_parallel_optimizer
 from mindspore.context import ParallelMode
 from mindspore import _checkparam as validator
 from mindspore import ops, nn
@@ -31,15 +31,12 @@ from mindspore.common import dtype as mstype
 from mindspore.common.parameter import Parameter, ParameterTuple
 from mindspore.common.tensor import Tensor
 from mindspore.ops.primitive import _primexpr
-from mindspore.ops import composite as C
-from mindspore.ops import functional as F
-from mindspore.ops import operations as P
 from mindspore.ops.operations.comm_ops import _VirtualDataset
 from mindspore.nn.cell import Cell
 from mindspore.nn.wrap.grad_reducer import DistributedGradReducer
 from mindspore.utils import ExitByRequest
 
-_get_datatype = C.MultitypeFuncGraph("_get_datatype")
+_get_datatype = ops.MultitypeFuncGraph("_get_datatype")
 
 
 @_get_datatype.register("Tensor")
@@ -53,10 +50,10 @@ def _tensors_get_datatype(param):
     Returns:
         mstype, the datatype of parameter.
     """
-    return F.dtype(param)
+    return ops.dtype(param)
 
 
-_cast_datatype = C.MultitypeFuncGraph("_cast_datatype")
+_cast_datatype = ops.MultitypeFuncGraph("_cast_datatype")
 
 
 @_cast_datatype.register("TypeType", "Tensor")
@@ -71,7 +68,7 @@ def _tensors_cast_datatype(datatype, param):
     Returns:
         Tensor, the parameter after operation.
     """
-    return F.cast(param, datatype)
+    return ops.cast(param, datatype)
 
 
 class WithLossCell(Cell):
@@ -195,7 +192,7 @@ class WithGradCell(Cell):
         self.network = network
         self.loss_fn = loss_fn
         self.weights = ParameterTuple(network.trainable_params())
-        self.grad = C.GradOperation(get_by_list=True, sens_param=(sens is not None))
+        self.grad = ops.GradOperation(get_by_list=True, sens_param=(sens is not None))
         self.sens = sens
         if loss_fn is None:
             self.network_with_loss = network
@@ -303,7 +300,7 @@ class ForwardValueAndGrad(Cell):
         self.get_all = get_all
         self.get_by_list = get_by_list
         self.sens_param = sens_param
-        self.grad = C.GradOperation(get_all=self.get_all, get_by_list=self.get_by_list, sens_param=self.sens_param)
+        self.grad = ops.GradOperation(get_all=self.get_all, get_by_list=self.get_by_list, sens_param=self.sens_param)
         self._get_attr_from_cell(network)
 
     def construct(self, *inputs):
@@ -385,8 +382,8 @@ class TrainOneStepCell(Cell):
         self.network.set_grad()
         self.optimizer = optimizer
         self.weights = self.optimizer.parameters
-        self.grad = C.GradOperation(get_by_list=True, sens_param=True)
-        self.grad_no_sens = C.GradOperation(get_by_list=True)
+        self.grad = ops.GradOperation(get_by_list=True, sens_param=True)
+        self.grad_no_sens = ops.GradOperation(get_by_list=True)
         self.sens = sens
         if self.sens == 0:
             raise ValueError("The input argument of 'sens' can not be 0.")
@@ -400,8 +397,7 @@ class TrainOneStepCell(Cell):
         self.reducer_flag = False
         self.grad_reducer = nn.Identity()
         self.parallel_mode = _get_parallel_mode()
-        self.reducer_flag = self.parallel_mode in (ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL) or \
-                            _is_pynative_parallel()
+        self.reducer_flag = self.parallel_mode in (ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL)
         if self.reducer_flag:
             self.mean = _get_gradients_mean()
             self.degree = _get_device_num()
@@ -428,12 +424,12 @@ class TrainOneStepCell(Cell):
         if not self.sense_flag:
             return self._no_sens_impl(*inputs)
         loss = self.network(*inputs)
-        sens = F.fill(loss.dtype, loss.shape, self.sens)
+        sens = ops.fill(loss.dtype, loss.shape, self.sens)
         grads = self.grad(self.network, self.weights)(*inputs, sens)
         grads = self.grad_reducer(grads)
         if self.use_graceful_exit:
             grads = self.graceful_exit.exit_by_request(grads, self.init_param, self.exit_param)
-        loss = F.depend(loss, self.optimizer(grads))
+        loss = ops.depend(loss, self.optimizer(grads))
         if self.return_grad:
             grad_with_param_name = {}
             for index, value in enumerate(grads):
@@ -448,7 +444,7 @@ class TrainOneStepCell(Cell):
         grads = self.grad_reducer(grads)
         if self.use_graceful_exit:
             grads = self.graceful_exit.exit_by_request(grads, self.init_param, self.exit_param)
-        loss = F.depend(loss, self.optimizer(grads))
+        loss = ops.depend(loss, self.optimizer(grads))
         if self.return_grad:
             grad_with_param_name = {}
             for index, value in enumerate(grads):
@@ -496,7 +492,7 @@ class GetNextSingleOp(Cell):
 
     def __init__(self, dataset_types, dataset_shapes, queue_name):
         super(GetNextSingleOp, self).__init__()
-        self.get_next = P.GetNext(dataset_types, dataset_shapes, len(dataset_types), queue_name)
+        self.get_next = ops.GetNext(dataset_types, dataset_shapes, len(dataset_types), queue_name)
 
     def construct(self):
         return self.get_next()
@@ -533,22 +529,22 @@ class _VirtualDatasetCell(Cell):
 
 
 def _pipeline_clear_grad(accu_grad, grad):
-    accu_grad = F.depend(accu_grad, grad)
-    zeros = F.zeros_like(accu_grad)
-    return F.assign(accu_grad, zeros)
+    accu_grad = ops.depend(accu_grad, grad)
+    zeros = ops.zeros_like(accu_grad)
+    return ops.assign(accu_grad, zeros)
 
 def grad_scale(scale, grad):
     """grad_scale"""
     new_grad = scale * grad
     grad = ops.depend(grad, new_grad)
-    zeros = F.zeros_like(grad)
-    new_grad = ops.depend(new_grad, F.assign(grad, zeros))
+    zeros = ops.zeros_like(grad)
+    new_grad = ops.depend(new_grad, ops.assign(grad, zeros))
     return new_grad
 
 
 @_primexpr
 def _check_shape_value_on_axis_divided_by_target_value(input_shape, micro_size):
-    if F.isconstant(input_shape[0]) is False:
+    if ops.isconstant(input_shape[0]) is False:
         return
     if input_shape[0] % micro_size != 0:
         raise ValueError(f"For micro batch initialization, the 0th dimension shape of input({input_shape[0]}) must be "
@@ -564,9 +560,9 @@ class _MicroBatch(Cell):
     """
     def __init__(self, micro_size):
         super(_MicroBatch, self).__init__()
-        self.shape = P.Shape()
+        self.shape = ops.Shape()
         self.micro_size = micro_size
-        self.strided_slice = P.StridedSlice()
+        self.strided_slice = ops.StridedSlice()
 
     def construct(self, i, *inputs):
         """construct for _MicroBatch."""
@@ -628,7 +624,7 @@ class GradAccumulationCell(Cell):
             micro_input = _MicroBatch(micro_size)
             micro_input.strided_slice.add_prim_attr("grad_accu_num", micro_size)
             self.micro_inputs.append(micro_input)
-            self.add = P.Add().add_prim_attr("forward_end", i)
+            self.add = ops.Add().add_prim_attr("forward_end", i)
             self.add_list.append(self.add)
         self._get_attr_from_cell(network)
 
@@ -685,10 +681,10 @@ class _TrainGradAccuStepCell(TrainOneStepCell):
         grads = self.grad_no_sens(self.network, self.weights)(*inputs)
         accu_grads = ops.depend(self.accu_grads, grads)
         if self.opt_shard:
-            grads = self.hyper_map(F.partial(grad_scale, self.sens), grads)
+            grads = self.hyper_map(ops.partial(grad_scale, self.sens), grads)
             succ = self.optimizer(grads)
         else:
-            accu_grads = self.hyper_map(F.partial(grad_scale, self.sens), accu_grads)
+            accu_grads = self.hyper_map(ops.partial(grad_scale, self.sens), accu_grads)
             succ = self.optimizer(accu_grads)
         loss = ops.depend(loss, succ)
         clear = self.hyper_map(_pipeline_clear_grad, accu_grads, grads)
@@ -794,8 +790,8 @@ class WithEvalCell(Cell):
     def construct(self, data, label):
         outputs = self._network(data)
         if self.add_cast_fp32:
-            label = F.mixed_precision_cast(mstype.float32, label)
-            outputs = F.cast(outputs, mstype.float32)
+            label = ops.mixed_precision_cast(mstype.float32, label)
+            outputs = ops.cast(outputs, mstype.float32)
         loss = self._loss_fn(outputs, label)
         return loss, outputs, label
 
@@ -845,7 +841,7 @@ class ParameterUpdate(Cell):
         self._param = param
 
     def construct(self, x):
-        F.assign(self._param, x)
+        ops.assign(self._param, x)
         return x
 
 
@@ -861,21 +857,21 @@ class _BroadCastCell(Cell):
         super(_BroadCastCell, self).__init__()
         from mindspore.communication.management import get_group_size, create_group
         from mindspore import context
-        self.map_ = C.Map()
+        self.map_ = ops.Map()
         self.params = tuple(params)
-        if context.get_context("device_target") == "Ascend" and context.get_context("mode") != context.PYNATIVE_MODE:
+        if context.get_context("device_target") == "Ascend":
             rank_list = [id for id in range(0, get_group_size())]
             create_group("BroadcastWorldGroup", rank_list)
-            self.broadcast = P.Broadcast(0, group="BroadcastWorldGroup")
+            self.broadcast = ops.Broadcast(0, group="BroadcastWorldGroup")
         else:
-            self.broadcast = P.Broadcast(0)
+            self.broadcast = ops.Broadcast(0)
         self.add_flags(skip_auto_parallel_compile=True)
 
     def construct(self):
-        datatypes = self.map_(F.partial(_get_datatype), self.params)
-        params = self.map_(F.partial(_cast_datatype, mstype.float32), self.params)
+        datatypes = self.map_(ops.partial(_get_datatype), self.params)
+        params = self.map_(ops.partial(_cast_datatype, mstype.float32), self.params)
         params = self.broadcast(params)
-        new_params = self.map_(F.partial(_cast_datatype), datatypes, params)
+        new_params = self.map_(ops.partial(_cast_datatype), datatypes, params)
         return new_params
 
 
@@ -892,6 +888,8 @@ class PipelineCell(Cell):
         micro_size (int): MicroBatch size.
         stage_config (dict, optional): The stage configuration for each cell's execution in pipeline parallel.
             Default ``None``.
+        segment_config (dict, optional): The segment configuration for each cell's execution in pipeline parallel.
+            Default ``None``.
 
     Supported Platforms:
         ``Ascend`` ``GPU``
@@ -903,7 +901,7 @@ class PipelineCell(Cell):
         >>> net = LeNet5()
         >>> net = nn.PipelineCell(net, 4)
     """
-    def __init__(self, network, micro_size, stage_config=None):
+    def __init__(self, network, micro_size, stage_config=None, segment_config=None):
         super(PipelineCell, self).__init__(auto_prefix=False)
         self.network = network
         self.micro_inputs = nn.CellList()
@@ -921,7 +919,7 @@ class PipelineCell(Cell):
         for i in range(micro_size):
             micro_input = _MicroBatch(micro_size)
             self.micro_inputs.append(micro_input)
-            self.add = P.Add().add_prim_attr("pipeline_end", i)
+            self.add = ops.Add().add_prim_attr("pipeline_end", i)
             self.add_list.append(self.add)
         self._get_attr_from_cell(network)
 
@@ -959,6 +957,39 @@ class PipelineCell(Cell):
                     print(cell_name)
                 raise KeyError("For 'PipelineCell', the argument 'stage_config' : {} is not "
                                "found in 'network' : {}".format(config_dict, network))
+        if segment_config is None:
+            return
+        self._config_segment(segment_config)
+
+
+    def _config_segment(self, segment_config=None):
+        """
+        Config segment num for cell.
+        """
+        config_dict = segment_config.copy()
+        for cell_name, cell in self.network.cells_and_names():
+            if cell_name in segment_config:
+                setattr(cell, "pipeline_segment", segment_config[cell_name])
+                del config_dict[cell_name]
+        if str(self.network) in segment_config:
+            setattr(self.network, "pipeline_segment", segment_config[str(self.network)])
+            del config_dict[str(self.network)]
+        # if there are any config elements left, print them
+        if config_dict:
+            for config_cell_name, config_segment_num in config_dict.items():
+                logger.error("pipeline_cell segment_config set pipeline_segment fail!")
+                logger.warning("config cell name:" + str(config_cell_name) +
+                               " config segment num:" + str(config_segment_num))
+            logger.warning("network:" + str(self.network))
+            logger.warning("cell name available:")
+            for cell_name, _ in self.network.cells_and_names():
+                logger.warning(cell_name)
+            raise KeyError("For 'PipelineCell', the argument 'segment_config' : {} is not "
+                           "found in 'network' : {}".format(config_dict, self.network))
+
+
+    def shard(self, in_strategy, out_strategy=None, parameter_plan=None, device="Ascend", level=0):
+        raise ValueError("For 'PipelineCell', no 'shard' on 'PipelineCell' is allowed.")
 
     def construct(self, *inputs):
         ret = None
@@ -1011,7 +1042,7 @@ class MicroBatchInterleaved(Cell):
         self.network = network
         self.interleave_num = interleave_num
         self.interleave_inputs = nn.CellList()
-        self.add = P.Add().add_prim_attr("micro_interleaved_add_flag", True)
+        self.add = ops.Add().add_prim_attr("micro_interleaved_add_flag", True)
         for _ in range(interleave_num):
             interleave_data = _MicroBatch(interleave_num)
             interleave_data.strided_slice.add_prim_attr("strided_slice_flag", True)

mindspore/nn/wrap/grad_reducer.py

@@ -20,7 +20,7 @@ from mindspore import log as logger
 from mindspore.nn.cell import Cell
 from mindspore.communication.management import GlobalComm, get_group_size
 from mindspore.common.sparse_tensor import RowTensorInner
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 from mindspore.ops.operations.comm_ops import AllReduce, AllGather
 from mindspore.parallel._auto_parallel_context import auto_parallel_context
 import mindspore.common.dtype as mstype
@@ -33,7 +33,7 @@ from mindspore.parallel._utils import _get_enable_parallel_optimizer
 __all__ = ['DistributedGradReducer']
 
 
-reduce_opt = C.MultitypeFuncGraph("reduce_opt")
+reduce_opt = ops.MultitypeFuncGraph("reduce_opt")
 
 
 def _init_allreduce_operators(length, split_indices, group=GlobalComm.WORLD_COMM_GROUP):
@@ -114,7 +114,7 @@ def _tensors_allreduce(degree, mean, allgather, allreduce, allreduce_filter, gra
     if allreduce_filter:
         grad = allreduce(grad)
         if mean:
-            grad = F.tensor_mul(grad, F.cast(degree, F.dtype(grad)))
+            grad = ops.tensor_mul(grad, ops.cast(degree, ops.dtype(grad)))
         return grad
     return grad
 
@@ -135,7 +135,7 @@ def _tensors_allreduce_post(degree, mean, allreduce_filter, grad):
     """
     if allreduce_filter:
         if mean:
-            grad = F.tensor_mul(grad, F.cast(degree, F.dtype(grad)))
+            grad = ops.tensor_mul(grad, ops.cast(degree, ops.dtype(grad)))
             return grad
     return grad
 
@@ -163,7 +163,7 @@ def _tensors_allreduce_ps(degree, mean, allgather, allreduce, allreduce_filter,
     if allreduce_filter:
         grad = allreduce(grad)
         if mean:
-            grad = F.tensor_mul(grad, F.cast(degree, F.dtype(grad)))
+            grad = ops.tensor_mul(grad, ops.cast(degree, ops.dtype(grad)))
         return grad
     return grad
 
@@ -189,7 +189,7 @@ def _tensors_allreduce_with_sparse(degree, mean, allgather, allreduce, allreduce
         indices = allgather(grad.indices)
         dout = allgather(grad.values)
         if mean:
-            dout = F.tensor_mul(dout, F.cast(degree, F.dtype(dout)))
+            dout = ops.tensor_mul(dout, ops.cast(degree, ops.dtype(dout)))
         grad = RowTensorInner(indices, dout, grad.dense_shape)
     return grad
 
@@ -219,12 +219,12 @@ def _tensors_allreduce_with_sparse_ps(degree, mean, allgather, allreduce, allred
         indices = allgather(grad.indices)
         dout = allgather(grad.values)
         if mean:
-            dout = F.tensor_mul(dout, F.cast(degree, F.dtype(dout)))
+            dout = ops.tensor_mul(dout, ops.cast(degree, ops.dtype(dout)))
         grad = RowTensorInner(indices, dout, grad.dense_shape)
     return grad
 
 
-_get_datatype = C.MultitypeFuncGraph("_get_datatype")
+_get_datatype = ops.MultitypeFuncGraph("_get_datatype")
 
 
 @_get_datatype.register("Tensor")
@@ -238,7 +238,7 @@ def _tensors_get_datatype(grad):
     Returns:
         mstype, the datatype of gradient.
     """
-    return F.dtype(grad)
+    return ops.dtype(grad)
 
 
 @_get_datatype.register("RowTensor")
@@ -252,10 +252,10 @@ def _tensors_get_datatype_with_sparse(grad):
     Returns:
         mstype, the datatype of gradient.
     """
-    return F.dtype(grad.values)
+    return ops.dtype(grad.values)
 
 
-_cast_datatype = C.MultitypeFuncGraph("_cast_datatype")
+_cast_datatype = ops.MultitypeFuncGraph("_cast_datatype")
 
 
 @_cast_datatype.register("TypeType", "Tensor")
@@ -270,7 +270,7 @@ def _tensors_cast_datatype(datatype, grad):
     Returns:
         Tensor, the gradient tensor after operation.
     """
-    return F.cast(grad, datatype)
+    return ops.cast(grad, datatype)
 
 
 @_cast_datatype.register("TypeType", "RowTensor")
@@ -285,7 +285,7 @@ def _tensors_cast_datatype_with_sparse(datatype, grad):
     Returns:
         RowTensor, the gradient after operation.
     """
-    dout = F.cast(grad.values, datatype)
+    dout = ops.cast(grad.values, datatype)
     return RowTensorInner(grad.indices, dout, grad.dense_shape)
 
 
@@ -361,7 +361,7 @@ class DistributedGradReducer(Cell):
         ...     def construct(self, *args):
         ...         weights = self.weights
         ...         loss = self.network(*args)
-        ...         sens = F.fill(ops.DType()(loss), ops.Shape()(loss), self.sens)
+        ...         sens = ops.fill(ops.DType()(loss), ops.Shape()(loss), self.sens)
         ...         grads = self.grad(self.network, weights)(*args, sens)
         ...         if self.reducer_flag:
         ...             # apply grad reducer on grads
@@ -394,8 +394,7 @@ class DistributedGradReducer(Cell):
 
     def __init__(self, parameters, mean=None, degree=None, fusion_type=1, group=GlobalComm.WORLD_COMM_GROUP):
         super(DistributedGradReducer, self).__init__(auto_prefix=False)
-        self._check_parallel_mode()
-        self.map_ = C.Map()
+        self.map_ = ops.Map()
         self.mean = mean
         if mean is None:
             self.mean = auto_parallel_context().get_gradients_mean()
@@ -443,54 +442,47 @@ class DistributedGradReducer(Cell):
         Returns:
             new_grads (Union[Tensor, tuple[Tensor]]), the gradient tensor or tuple after operation.
         """
-        datatypes = self.map_(F.partial(_get_datatype), grads)
-        grads = self.map_(F.partial(_cast_datatype, mstype.float32), grads)
+        datatypes = self.map_(ops.partial(_get_datatype), grads)
+        grads = self.map_(ops.partial(_cast_datatype, mstype.float32), grads)
 
         if self.split_fusion:
             if self.enable_parameter_server:
-                new_grad = self.map_(F.partial(reduce_opt, self.degree, self.mean, self.allgather),
+                new_grad = self.map_(ops.partial(reduce_opt, self.degree, self.mean, self.allgather),
                                      self.op_list, self.allreduce_filter, grads, self.ps_parameters)
             else:
-                new_grad = self.map_(F.partial(reduce_opt, self.degree, self.mean, self.allgather),
+                new_grad = self.map_(ops.partial(reduce_opt, self.degree, self.mean, self.allgather),
                                      self.op_list, self.allreduce_filter, grads)
         else:
             if self.enable_parameter_server:
-                new_grad = self.map_(F.partial(reduce_opt, self.degree, self.mean, self.allgather,
-                                               self.allreduce), self.allreduce_filter, grads, self.ps_parameters)
+                new_grad = self.map_(ops.partial(reduce_opt, self.degree, self.mean, self.allgather,
+                                                 self.allreduce), self.allreduce_filter, grads, self.ps_parameters)
             else:
-                new_grad = self.map_(F.partial(reduce_opt, self.degree, self.mean, self.allgather,
-                                               self.allreduce), self.allreduce_filter, grads)
-        new_grad = self.map_(F.partial(_cast_datatype), datatypes, new_grad)
+                new_grad = self.map_(ops.partial(reduce_opt, self.degree, self.mean, self.allgather,
+                                                 self.allreduce), self.allreduce_filter, grads)
+        new_grad = self.map_(ops.partial(_cast_datatype), datatypes, new_grad)
         return new_grad
 
-    def _check_parallel_mode(self):
-        """check parallel mode"""
-        parallel_mode = context.get_auto_parallel_context('parallel_mode')
-        if context.get_context('mode') == context.GRAPH_MODE and parallel_mode in (
-                context.ParallelMode.SEMI_AUTO_PARALLEL, context.ParallelMode.AUTO_PARALLEL):
-            raise RuntimeError("{} can not use DistributedGradReducer in graph mode".format(parallel_mode))
 
-
-grad_scale = C.MultitypeFuncGraph("grad_scale")
-shard_grad_scale = C.MultitypeFuncGraph("shard_grad_scale")
-reciprocal = P.Reciprocal()
+grad_scale = ops.MultitypeFuncGraph("grad_scale")
+shard_grad_scale = ops.MultitypeFuncGraph("shard_grad_scale")
+reciprocal = ops.Reciprocal()
 
 
 @grad_scale.register("Tensor", "Tensor", "Tensor")
 def tensor_grad_scale_pipeline(scale, grad, accu_grad):
-    accu_grad = F.depend(accu_grad, grad)
+    accu_grad = ops.depend(accu_grad, grad)
     new_grad = accu_grad * reciprocal(scale)
-    accu_grad = F.depend(accu_grad, new_grad)
-    zeros = F.tensor_mul(accu_grad, 0.0)
-    new_grad = F.depend(new_grad, F.assign(accu_grad, zeros))
+    accu_grad = ops.depend(accu_grad, new_grad)
+    zeros = ops.tensor_mul(accu_grad, 0.0)
+    new_grad = ops.depend(new_grad, ops.assign(accu_grad, zeros))
     return new_grad
 
 
 @shard_grad_scale.register("Tensor", "Tensor", "Tensor")
 def tensor_shard_grad_scale_pipeline(scale, grad, accu_grad):
     new_grad = grad * reciprocal(scale)
-    accu_grad = F.depend(accu_grad, new_grad)
-    new_grad = F.depend(new_grad, F.assign(accu_grad, F.zeros_like(accu_grad)))
+    accu_grad = ops.depend(accu_grad, new_grad)
+    new_grad = ops.depend(new_grad, ops.assign(accu_grad, ops.zeros_like(accu_grad)))
     return new_grad
 
 
@@ -563,7 +555,7 @@ class PipelineGradReducer(Cell):
         >>> net.layer3.pipeline_stage = 1
         >>> loss_fn = nn.CrossEntropyLoss()
         >>> optimizer = nn.SGD(net.trainable_params(), 1e-2)
-        >>> net_with_loss = nn.Pipeline(nn.WithLossCell(net, loss_fn), 2)
+        >>> net_with_loss = nn.PipelineCell(nn.WithLossCell(net, loss_fn), 2)
         >>> net_with_loss.set_train()
         >>> def forward_fn(inputs, target):
         ...     loss = net_with_loss(inputs, target)
@@ -587,12 +579,11 @@ class PipelineGradReducer(Cell):
     """
     def __init__(self, parameters, scale_sense=1.0, opt_shard=None):
         super(PipelineGradReducer, self).__init__(auto_prefix=False)
-        self._check_mode()
         self.accu_grads = parameters.clone(prefix="accu_grads", init="zeros")
         self.grad_reducer = Identity()
         self.degree = Tensor(1, mstype.float32)
         self.scale_sense = Parameter(scale_sense, name='scale_sense')
-        self.hyper_map = C.HyperMap()
+        self.hyper_map = ops.HyperMap()
         if opt_shard is None:
             self.opt_shard = _get_enable_parallel_optimizer()
         else:
@@ -603,15 +594,9 @@ class PipelineGradReducer(Cell):
         new_grads = None
         if self.opt_shard:
             grads = self.grad_reducer(grads)
-            new_grads = self.hyper_map(F.partial(shard_grad_scale, self.scale_sense * self.degree),
+            new_grads = self.hyper_map(ops.partial(shard_grad_scale, self.scale_sense * self.degree),
                                        grads, self.accu_grads)
         else:
             accu_grads = self.grad_reducer(self.accu_grads)
-            new_grads = self.hyper_map(F.partial(grad_scale, self.scale_sense * self.degree), grads, accu_grads)
+            new_grads = self.hyper_map(ops.partial(grad_scale, self.scale_sense * self.degree), grads, accu_grads)
         return new_grads
-
-    def _check_mode(self):
-        """check parallel mode"""
-        mode = context.get_context('mode')
-        if mode != context.GRAPH_MODE:
-            raise RuntimeError(f"PipelineGradReducer only support graph mode, but get {mode}")