mindspore 2.1.0__cp39-cp39-win_amd64.whl → 2.2.10__cp39-cp39-win_amd64.whl

mindspore/ops/operations/other_ops.py

@@ -19,7 +19,6 @@ from mindspore import log as logger
 from mindspore.ops import signature as sig
 from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
-from mindspore.common._decorator import deprecated
 from mindspore.ops.primitive import Primitive, PrimitiveWithCheck, PrimitiveWithInfer, prim_attr_register
 from mindspore.ops.operations._pyfunc_registry import add_pyfunc
 from mindspore._c_expression import typing
@@ -738,27 +737,6 @@ class Pull(PrimitiveWithInfer):
         return mstype.float32
 
 
-class identity(Primitive):
-    """
-    The :class:`mindspore.ops.identity` interface is deprecated, please use the :class:`mindspore.nn.Identity` instead.
-
-    Supported Platforms:
-        Deprecated
-    """
-
-    # Side effect will propagated from the first argument to return value.
-    side_effect_propagate = 1
-
-    @prim_attr_register
-    def __init__(self):
-        """Initialize identity."""
-        self.add_prim_attr('side_effect_propagate', 1)
-
-    @deprecated('2.0', 'nn.Identity', False)
-    def __call__(self, x):
-        return x
-
-
 class PyInterpret(Primitive):
     r"""
     Interpret Python expression.

mindspore/ops/operations/random_ops.py

@@ -83,15 +83,10 @@ class TruncatedNormal(Primitive):
     Note:
         - The value of `shape` must be greater than zero. The output length can not exceed 1000000.
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -152,15 +147,10 @@ class StandardNormal(Primitive):
 
     Note:
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -208,15 +198,10 @@ class StandardLaplace(Primitive):
 
     Note:
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -266,15 +251,10 @@ class RandomGamma(Primitive):
 
     Note:
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -380,15 +360,10 @@ class Gamma(PrimitiveWithInfer):
 
     Note:
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -468,15 +443,10 @@ class ParameterizedTruncatedNormal(Primitive):
     Note:
         - The value in tensor `min` must be strictly less than `max` at any position after broadcasting.
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -551,15 +521,10 @@ class Poisson(PrimitiveWithInfer):
 
     Note:
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -630,15 +595,10 @@ class RandomPoisson(Primitive):
 
     Note:
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -705,15 +665,10 @@ class UniformInt(Primitive):
     Note:
         - The number in tensor minval must be strictly less than maxval at any position after broadcasting.
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -769,15 +724,16 @@ class UniformReal(Primitive):
 
     Note:
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
+        - Currently, on the Ascend platform, `shape` as a Tensor is not supported.
+          This is supported on CPU/GPU platforms. When the input is a Tensor,
+          the supported data types are as follows:
+
+          - GPU: int32, int64.
+          - CPU: int16, int32, int64.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -787,7 +743,6 @@ class UniformReal(Primitive):
 
     Inputs:
         - **shape** (Union[tuple, Tensor]) - The shape of tensor to be generated. Only constant value is allowed.
-          Supported dtypes: int16, int32, int64.
 
     Outputs:
         Tensor. The shape that the input 'shape' denotes. The dtype is float32.
@@ -809,6 +764,7 @@ class UniformReal(Primitive):
         >>> print(result)
         (2, 2)
     """
+
     @prim_attr_register
     def __init__(self, seed=0, seed2=0):
         """Initialize UniformReal"""
@@ -826,15 +782,10 @@ class RandomChoiceWithMask(Primitive):
 
     Note:
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         count (int, optional): Number of items expected to get and the number must be greater than 0. Default: ``256`` .
@@ -850,8 +801,8 @@ class RandomChoiceWithMask(Primitive):
     Outputs:
         Two tensors, the first one is the index tensor and the other one is the mask tensor.
 
-        - **index** (Tensor) - The output shape is 2-D.
-        - **mask** (Tensor) - The output shape is 1-D.
+        - **index** (Tensor) - The output shape is 2-D, its shape is :math:`(count, rank of input_x)`.
+        - **mask** (Tensor) - The output shape is 1-D, its shape is :math:`(count)`.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -945,15 +896,10 @@ class Multinomial(Primitive):
         - The rows of input do not need to sum to one (in which case we use the values as weights),
           but must be non-negative, finite and have a non-zero sum.
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,
@@ -1024,8 +970,8 @@ class MultinomialWithReplacement(Primitive):
     Inputs:
         - **x** (Tensor) - the input tensor containing the cumsum of probabilities, must be 1 or 2
           dimensions.
-        - **seed** (Tensor) - If `seed` is set to -1, and `offset` is set to 0, the random number
-          generator is seeded by a random seed. Otherwise, it is seeded by the given seed.
+        - **seed** (Tensor) - If `seed` and 'offset' are both set to 0, the random number generator
+          is seeded by a random seed. Otherwise, it is seeded by the given seed and offset.
           Supported dtype: int64.
         - **offset** (Tensor) - Offset used to avoid seed collision. Supported dtype: int64.
 
@@ -1072,7 +1018,9 @@ class UniformCandidateSampler(Primitive):
         range_max (int): The number of possible classes, must be non-negative.
         seed (int, optional): Used for random number generation, must be non-negative. If seed has a value of 0,
             the seed will be replaced with a randomly generated value. Default: ``0`` .
-        remove_accidental_hits (bool, optional): Whether accidental hit is removed. Default: ``False`` .
+        remove_accidental_hits (bool, optional): Whether accidental hit is removed.
+            Accidental hit is when one of the true classes matches one of the sample classes.
+            Set ``True`` to remove which accidentally sampling the true class as sample class. Default: ``False`` .
 
     Inputs:
         - **true_classes** (Tensor) - A Tensor. The target classes with a Tensor shape of
@@ -1128,7 +1076,6 @@ class UniformCandidateSampler(Primitive):
         self.add_prim_attr("side_effect_hidden", True)
 
 
-
 class LogUniformCandidateSampler(Primitive):
     r"""
     Generates random labels with a log-uniform distribution for sampled_candidates.
@@ -1206,15 +1153,10 @@ class RandomShuffle(Primitive):
 
     Note:
         - Random seed: a set of regular random numbers can be obtained through some complex mathematical algorithms,
-          and the random seed is the initial value of this random number. If the random seed is the same in two
+          and the random seed determines the initial value of this random number. If the random seed is the same in two
           separate calls, the random number generated will not change.
-        - Global random seed and operator-level random seed are not set or both set to 0: behavior is completely random.
-        - Global random seed is set, but operator-level random seed is not set: A global random seed will splice
-          with 0 to generate random number.
-        - Global random seed is not set, operator-level random seed is set: 0
-          splices with the operator-level random seed to generate random number.
-        - Both Global random and operator-level random seed are set: the global random seed will splice with the
-          operator-level random seed to generate random number.
+        - Using the Philox algorithm to scramble seed and seed2 to obtain random seed so that the user doesn't need
+          to worry about which seed is more important.
 
     Args:
         seed (int, optional): The operator-level random seed, used to generate random numbers,

mindspore/ops/operations/sparse_ops.py

@@ -1615,7 +1615,7 @@ class SparseMatrixSoftmax(Primitive):
         if not isinstance(dtype, (type(mstype.float32), type(mstype.single), type(mstype.float64),
                                   type(mstype.double))):
             raise TypeError(
-                "Only float32 and float64 type data are supported, but got {}".format(dtype))
+                f"Only float32 and float64 type data are supported, but got {dtype}")
         self.add_prim_attr("dtype", dtype)
         self.init_prim_io_names(inputs=['x_dense_shape', 'x_batch_pointers', 'x_row_pointers',
                                         'x_col_indices', 'x_values'],
@@ -2602,6 +2602,8 @@ class RaggedTensorToTensor(Primitive):
                 raise ValueError(
                     f"For {self.name}, the each element of row_partition_types must be 'ROW_SPLITS' "
                     f"when row_splits tensor.")
+        self.num_row_partition_tensors = len(row_partition_types)
+        self.add_prim_attr("num_row_partition_tensors", self.num_row_partition_tensors)
 
 
 class SparseCross(Primitive):

mindspore/ops/primitive.py

@@ -25,7 +25,7 @@ from mindspore.parallel._ps_context import _is_ps_mode, _is_role_sched
 from mindspore.common.parameter import Parameter
 from mindspore.common.api import _pynative_executor
 from mindspore.common._stub_tensor import _convert_stub
-from mindspore._c_expression import Primitive_, prim_type
+from mindspore._c_expression import Primitive_, prim_type, typing
 from mindspore import _checkparam as Validator
 from mindspore.ops import signature as sig
 
@@ -746,6 +746,8 @@ def _check_contains_variable(item_dtype, item_value):
             if _check_contains_variable(item_dtype[i], element):
                 return True
     elif isinstance(item_value, dict):
+        if isinstance(item_dtype, typing.Keyword):
+            return item_value is None
         for i in range(len(item_value)):
             if _check_contains_variable(item_dtype[i], list(item_value.keys())[i]):
                 return True
@@ -756,9 +758,7 @@ def _check_contains_variable(item_dtype, item_value):
 
 
 def constexpr(fn=None, get_instance=True, name=None, reuse_result=True, check=True):
-    """
-    Creates a PrimitiveWithInfer operator that can infer the value at compile time. We can use it to define a function
-    to compute constant value using the constants in the constructor.
+    """Used to calculate constant in graph copmpiling process and improve compile performance in GRAPH_MODE.
 
     Args:
         fn (function): A `fn` use as the infer_value of the output operator. Default: ``None`` .
@@ -772,22 +772,27 @@ def constexpr(fn=None, get_instance=True, name=None, reuse_result=True, check=Tr
             and the warning message will raised if the parameter is not const value. Default: ``True`` .
 
     Examples:
-        >>> from mindspore.ops import constexpr
-        >>> a = (1, 2)
-        >>> # make an operator to calculate tuple len
-        >>> @constexpr
-        ... def tuple_len(x):
-        ...     return len(x)
+
+        >>> import mindspore as ms
+        >>> # define a constant calculate function with for loop inside and use use constexpr to accelerate the compile
+        >>> # process.
+        >>> @ms.constexpr
+        ... def for_loop_calculate(range_num):
+        ...     out = 0
+        ...     for i in range(range_num):
+        ...         if i %2 == 0 and i % 7 != 0:
+        ...             out = out + i
+        ...     return out // range_num
         ...
-        >>> print(tuple_len(a))
-        2
-        >>> # make an operator class to calculate tuple len
-        >>> @constexpr(get_instance=False, name="TupleLen")
-        ... def tuple_len_class(x):
-        ...     return len(x)
+        >>> # construct a net and run with GRAPH_MODE.
+        >>> @ms.jit
+        ... def my_func(x):
+        ...     new_shape = for_loop_calculate(100000)
+        ...     return ms.ops.broadcast_to(x, (new_shape, ))
         ...
-        >>> print(tuple_len_class()(a))
-        2
+        >>> out = my_func(ms.Tensor([1]))
+        >>> print(out.shape)
+        >>> (21428, )
     """
 
     def deco(fn):
@@ -844,6 +849,7 @@ def _primexpr(fn=None, get_instance=True, name=None, reuse_result=True):
         reuse_result (bool): If ``True`` , the operator will be executed once and reuse the result next time,
                              otherwise the operator will always be executed. Default: ``True`` .
     """
+
     def deco(fn):
         """Decorator for CompileOp."""
 

mindspore/parallel/_auto_parallel_context.py

@@ -62,6 +62,7 @@ class _ParallelOptimizerConfig:
     """
     GRADIENT_ACCUMULATION_SHARD = "gradient_accumulation_shard"
     PARALLEL_OPTIMIZER_THRESHOLD = "parallel_optimizer_threshold"
+    OPTIMIZER_WEIGHT_SHARD_SIZE = "optimizer_weight_shard_size"
 
 
 class _AutoParallelContext:
@@ -176,7 +177,6 @@ class _AutoParallelContext:
         if comm_type == _ParallelFusionConfig.REDUCESCATTER:
             self._context_handle.set_reducescatter_fusion_threshold_mb(fusion_threshold)
 
-
     def fusion_threshold_mb(self):
         """Get all reduce threshold."""
         self.check_context_handle()
@@ -229,6 +229,22 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_pipeline_stage_split_num()
 
+    def set_pipeline_segments(self, segments):
+        """Set the segments of the pipeline"""
+        if isinstance(segments, bool) or not isinstance(segments, int):
+            raise TypeError("For 'set_auto_parallel_context', the argument 'pipeline_segments' "
+                            "must be int, but got the type : {}.".format(type(segments)))
+        if segments < 1:
+            raise ValueError("For 'set_auto_parallel_context', the argument 'pipeline_segments' "
+                             "should be greater or equal 1, but got the value of segments : {}.".format(segments))
+        self.check_context_handle()
+        self._context_handle.set_pipeline_segment_split_num(segments)
+
+    def get_pipeline_segments(self):
+        """Get the stages of the pipeline"""
+        self.check_context_handle()
+        return self._context_handle.get_pipeline_segment_split_num()
+
     def set_gradients_mean(self, gradients_mean):
         """
         Set gradients_mean flag.
@@ -491,6 +507,9 @@ class _AutoParallelContext:
         Args:
             grad_accumulation_step (int): The grad accumulation step.
         """
+        if grad_accumulation_step > 1:
+            raise ValueError("The interface is deprecated. To use gradient accumulation, "
+                             "please use GradAccumulationCell in mindspore.nn.wrap.cell_wrapper.")
         self.check_context_handle()
         Validator.check_positive_int(grad_accumulation_step)
         self._context_handle.set_grad_accumulation_step(grad_accumulation_step)
@@ -758,6 +777,11 @@ class _AutoParallelContext:
                             .format(type(enable_parallel_optimizer)))
         self._context_handle.set_enable_parallel_optimizer(enable_parallel_optimizer)
 
+    def get_enable_fold_pipeline(self):
+        """Get parallel optimizer flag."""
+        self.check_context_handle()
+        return self._context_handle.get_enable_fold_pipeline()
+
     def get_enable_parallel_optimizer(self):
         """Get parallel optimizer flag."""
         self.check_context_handle()
@@ -767,8 +791,6 @@ class _AutoParallelContext:
         r"""
         Set the configure for parallel optimizer. The configure provides more detailed behavior control about parallel
         training when parallel optimizer is enabled.
-        Currently it supports the key `gradient_accumulation_shard`. The configure will be effective
-        when we use context.set_auto_parallel_context(enable_parallel_optimizer=True).
 
         Args:
             parallel_optimizer_config(dict): A dict contains the keys and values for setting the parallel optimizer
@@ -786,14 +808,21 @@ class _AutoParallelContext:
                                                  enabled, parameters with size smaller than this threshold will not be
                                                  sharded across the devices. Parameter size = shape[0] \* ... \*
                                                  shape[n] \* size(dtype). Non-negative. Unit: KB. Default: 64.
+            - optimizer_weight_shard_size(int): Set the optimizer weight shard group size if you want to specific the
+                                                maximum group size across devices when the parallel optimizer is
+                                                enabled. The numerical range can be (0, device_num]. Default value
+                                                is -1, which means the optimizer weight shard group size will
+                                                the data parallel group of each parameter. Default -1.
+
         """
         self.check_context_handle()
         grad_shard_name = _ParallelOptimizerConfig.GRADIENT_ACCUMULATION_SHARD
         threshold_name = _ParallelOptimizerConfig.PARALLEL_OPTIMIZER_THRESHOLD
+        optimizer_weight_shard_size_name = _ParallelOptimizerConfig.OPTIMIZER_WEIGHT_SHARD_SIZE
 
         for config_name in parallel_optimizer_config:
             unknown_config = []
-            if config_name not in [grad_shard_name, threshold_name]:
+            if config_name not in [grad_shard_name, threshold_name, optimizer_weight_shard_size_name]:
                 unknown_config.append(config_name)
 
             if unknown_config:
@@ -811,6 +840,11 @@ class _AutoParallelContext:
             self._context_handle.set_parallel_optimizer_threshold(
                 parallel_optimizer_config[threshold_name])
 
+        if optimizer_weight_shard_size_name in parallel_optimizer_config:
+            value = parallel_optimizer_config[optimizer_weight_shard_size_name]
+            Validator.check_positive_int(value)
+            self.set_optimizer_weight_shard_size(value)
+
     def get_grad_accumulation_shard(self):
         """Get grad accumulation shard."""
         self.check_context_handle()
@@ -890,6 +924,13 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_optimizer_weight_shard_size()
 
+    def set_ops_strategy_json_config(self, type, path, mode):
+        """
+         Set configuration of saving ops strategy in file .json.
+        """
+        self.check_context_handle()
+        self._context_handle.set_ops_strategy_json_config(type, path, mode)
+
     def set_optimizer_weight_shard_aggregated_save(self, optimizer_weight_shard_aggregated_save):
         """
         Set optimizer_weight_shard_aggregated_save.
@@ -1027,8 +1068,28 @@ class _AutoParallelContext:
             self.set_enable_all_gather_fusion(openstate)
             self.set_enable_reduce_scatter_fusion(openstate)
 
+def _set_ops_strategy_json_config(type="SAVE", path="", mode="all"):
+    """
+    Set strategy json configuration.
 
+    Args:
+        type (str): The parameter for choosing save or load .json file.
+        path (str): Path to save or load parallel strategy json.
+        mode (str): The parameter for choosing save all or important operators.
 
+    Raises:
+        KeyError: When type is not 'SAVE' or 'LOAD'.
+        KeyError: When mode is not 'all' or 'principal'.
+    """
+    dir_path = os.path.dirname(path)
+    if dir_path and not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+    check_type = ["SAVE", "LOAD"]
+    check_mode = ["all", "principal"]
+    if type in check_type and mode in check_mode:
+        auto_parallel_context().set_ops_strategy_json_config(type, path, mode)
+    else:
+        raise KeyError("Type must be 'SAVE' or 'LOAD' and mode must be 'all' or 'principal'")
 
 _AUTO_PARALLEL_CONTEXT = None
 
@@ -1053,6 +1114,7 @@ _set_auto_parallel_context_func_map = {
     "gradient_fp32_sync": auto_parallel_context().set_gradient_fp32_sync,
     "loss_repeated_mean": auto_parallel_context().set_loss_repeated_mean,
     "pipeline_stages": auto_parallel_context().set_pipeline_stages,
+    "pipeline_segments": auto_parallel_context().set_pipeline_segments,
     "parallel_mode": auto_parallel_context().set_parallel_mode,
     "search_mode": auto_parallel_context().set_strategy_search_mode,
     "auto_parallel_search_mode": auto_parallel_context().set_auto_parallel_search_mode,
@@ -1074,7 +1136,6 @@ _set_auto_parallel_context_func_map = {
     "strategy_ckpt_config": auto_parallel_context().set_strategy_ckpt_config,
     "comm_fusion": auto_parallel_context().set_comm_fusion}
 
-
 _get_auto_parallel_context_func_map = {
     "device_num": auto_parallel_context().get_device_num,
     "global_rank": auto_parallel_context().get_global_rank,
@@ -1111,7 +1172,6 @@ _get_auto_parallel_context_func_map = {
                  communi_parallel_mode=str, optimizer_weight_shard_size=int, sharding_propagation=bool,
                  optimizer_weight_shard_aggregated_save=bool, enable_alltoall=bool, comm_fusion=dict,
                  strategy_ckpt_config=dict)
-
 def _set_auto_parallel_context(**kwargs):
     """
     Set auto parallel context.
@@ -1247,8 +1307,8 @@ def _reset_auto_parallel_context():
     - strategy_ckpt_load_file: ""
     - strategy_ckpt_save_file: ""
     - enable_parallel_optimizer: False
-    - search_mode: dynamic_programming
-    - auto_parallel_search_mode: dynamic_programming
+    - search_mode: 'recursive_programming
+    - auto_parallel_search_mode: 'recursive_programming
     - sharding_propagation: False
     - pipeline_stages: 0
     - gradient_accumulation_shard: True

mindspore/parallel/_cost_model_context.py

@@ -475,7 +475,7 @@ class _CostModelContext:
         """
         if self._context_handle is None:
             raise ValueError("Context handle is none in context!!!")
-        return self._context_handle.rp_matmul_mem_coef()
+        return self._context_handle.get_rp_matmul_mem_coef()
 
     def set_costmodel_allreduce_fusion_computation_time_parameter(self, computation_time_parameter):
         """
@@ -693,7 +693,7 @@ def _set_rp_matmul_mem_coef(coef):
     cost_model_context().set_rp_matmul_mem_coef(coef)
 
 
-def _get_rp_matmul_mem_coef(self):
+def _get_rp_matmul_mem_coef():
     """
     Get the matmul memory coef which is used in the RP algorithm.
     """