mindspore 2.5.0__cp311-cp311-win_amd64.whl → 2.6.0__cp311-cp311-win_amd64.whl

mindspore/_extends/parse/standard_method.py

@@ -19,7 +19,7 @@
 from __future__ import absolute_import
 from mindspore import Tensor, CSRTensor, COOTensor, Parameter
 from mindspore import dtype as mstype
-from mindspore._c_expression import Tensor as Tensor_
+from mindspore._c_expression import TensorPy as Tensor_
 from mindspore.common import mutable
 from mindspore.common.generator import default_generator
 import mindspore.common._monad as monad
@@ -28,8 +28,9 @@ from mindspore.ops.composite.base import _append, _insert, _pop, _list_clear, _r
     _extend, _dict_setitem, _dict_clear, _haskey, _update, _fromkeys
 from mindspore.ops.operations._sequence_ops import TensorToTuple
 from mindspore.ops.auto_generate import trace_v2_op, inplace_addmm_op, inplace_index_put_op, inplace_normal_op, inplace_index_add_op
-from mindspore.ops.auto_generate import inplace_copy_op
+from mindspore.ops.auto_generate import inplace_copy_op, inplace_uniform_op, inplace_erfinv_op
 from mindspore.ops.auto_generate import inplace_scatter_add as inplace_scatter_add_
+from mindspore.ops.auto_generate import inplace_exponential_op
 
 from ... import _checkparam as validator
 from ..._checkparam import check_is_number, check_reshape_shp, check_axis_in_range, \
@@ -39,7 +40,7 @@ from ...ops import operations as P
 from ...ops import composite
 from ...ops.operations import array_ops
 from ...ops.composite import MultitypeFuncGraph, env_get, hyper_add, \
-    zeros_like, ones_like, repeat_elements, multitype_ops
+    zeros_like, ones_like, multitype_ops, _ones_like_for_grad
 from ...ops.composite.multitype_ops import _constexpr_utils as const_utils
 from ...ops.composite.multitype_ops import _compile_utils as compile_utils
 from ...ops.operations._inner_ops import Format
@@ -52,7 +53,7 @@ from ...ops.operations._sequence_ops import ListAppend, ListInsert, SequenceMax,
     SequenceIndex
 
 __all__ = ['MultitypeFuncGraph', 'env_get',
-           'hyper_add', 'zeros_like', 'ones_like']
+           'hyper_add', 'zeros_like', 'ones_like', '_ones_like_for_grad']
 
 shape_ = P.Shape()
 dtype_ = P.DType()
@@ -421,7 +422,7 @@ def hasattr(x, attr):  # pylint: disable=redefined-builtin
 
     Args:
         x (object): Input object.
-        attr (string): The name of attribute
+        attr (str): The name of attribute
 
     Returns:
         Boolean value, indicates whether the object x has attribute attr.
@@ -579,7 +580,7 @@ def transpose(x, *axis):
 
     Raises:
         TypeError: If input arguments have types not specified above.
-        ValueError: If the number of `axes` is not euqal to a.ndim.
+        ValueError: If the number of `axes` is not equal to a.ndim.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -1232,6 +1233,13 @@ def pow(x, y):  # pylint: disable=redefined-builtin
     return F.pow(x, y)
 
 
+def put_(x, index, source, accumulate=False):  # pylint: disable=redefined-builtin
+    """
+    Copies the elements from source into the positions specified by index.
+    """
+    return F.put_(x, index, source, accumulate)
+
+
 def log(x):
     """
     Calculate the logarithm of Tensor.
@@ -1274,12 +1282,12 @@ def logcumsumexp(input, axis):
     return F.logcumsumexp(input, axis)
 
 
-def logsumexp(input, axis, keepdims=False):
+def logsumexp(input, dim, keepdim=False):
     """
     Reduces a dimension of a tensor by calculating exponential for all elements in the dimension,
     then calculate logarithm of the sum.
     """
-    return F.logsumexp(input, axis, keepdims)
+    return F.logsumexp(input, dim, keepdim)
 
 
 def round_(x):
@@ -1289,12 +1297,11 @@ def round_(x):
     return F.round(x)
 
 
-def roll(x, shifts, dims):
+def roll(x, shifts, dims=None):
     """
     Rolls the elements of a tensor along an axis.
     """
-    dims = dims if dims is not None else 0
-    return F.Roll(shifts, dims)(x)
+    return F.roll(x, shifts, dims)
 
 
 def rot90(x, k, dims):
@@ -1362,11 +1369,11 @@ def remainder(input, divisor):
     return F.remainder(input, divisor)
 
 
-def unique_consecutive(input, return_idx=False, return_counts=False, axis=None):
+def unique_consecutive(input, return_inverse=False, return_counts=False, dim=None):
     """
     Returns the elements that are unique in each consecutive group of equivalent elements in the input tensor.
     """
-    return F.unique_consecutive(input, return_idx, return_counts, axis)
+    return F.unique_consecutive(input, return_inverse, return_counts, dim)
 
 
 def unique_with_pad(x, pad_num):
@@ -1850,7 +1857,7 @@ def searchsorted(x, v, side='left', sorter=None):
 
     if side not in ('left', 'right'):
         raise ValueError(f"For 'Tensor.searchsorted', the argument 'side' should be one of in "
-                            f"['left', 'right'], but got {side}.")
+                         f"['left', 'right'], but got {side}.")
     if not isinstance(v, Tensor):
         v = const_utils.make_tensor(v)
     if sorter is not None:
@@ -2177,71 +2184,53 @@ def nanmedian(input, axis=-1, keepdims=False):
     return F.nanmedian(input, axis, keepdims)
 
 
-def repeat(x, repeats, axis=None):
+def repeat(x, *args, repeats=None):
     """
     Repeat elements of an array.
 
     Args:
         x (Tensor): Input tensor.
-        repeats (Union[int, tuple, list]): The number of repetitions for each element.
-            `repeats` is broadcasted to fit the shape of the given axis.
-        axis (int, optional): The axis along which to repeat values. By default,
-            use the flattened input tensor, and return a flat output tensor.
+        args (*int): To simulate an overload like ``repeat(x, *repeats: int)``.
+        repeats (Union[int, tuple[int], list[int]]): The number of repetitions of `a` along
+            each axis. Requires that ``len(repeats) >= x.rank``.
 
     Returns:
-        Tensor, has the same shape as input tensor except along the given axis.
+        Tensor, the repeated output array.
 
     Raises:
-        ValueError: if axis is out of range.
         TypeError: if input is not a Tensor.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> import mindspore.numpy as np
-        >>> x = np.array(3)
-        >>> print(x.repeat(4))
-        [3 3 3 3]
-        >>> x = np.array([[1,2],[3,4]])
-        >>> print(x.repeat(2))
-        [1 1 2 2 3 3 4 4]
-        >>> print(x.repeat(3, axis=1))
-        [[1 1 1 2 2 2]
-        [3 3 3 4 4 4]]
-        >>> print(x.repeat([1,2], axis=0))
-        [[1 2]
-        [3 4]
-        [3 4]]
-    """
-    if not isinstance(repeats, (tuple, list)):
-        repeats = (repeats,)
-    for element in repeats:
-        if not isinstance(element, int):
-            const_utils.raise_type_error("Each element should be integer")
-    if axis is None:
-        x = ravel(x)
-        axis = 0
-    if not isinstance(axis, int):
-        const_utils.raise_type_error('axes should be integers')
-    check_axis_in_range(axis, x.ndim)
-    axis = axis + x.ndim if axis < 0 else axis
-
-    if len(repeats) == 1:
-        repeats = repeats[0]
-        if repeats == 0:
-            return empty_tensor(x.dtype)
-        return repeat_elements(x, repeats, axis)
-    size = x.shape[axis]
-    if len(repeats) != size:
-        const_utils.raise_value_error(
-            'operands could not be broadcast together')
-    subs = P.Split(axis, size)(x)
-    repeated_subs = []
-    for sub_item, rep in zip(subs, repeats):
-        if rep != 0:
-            repeated_subs.append(repeat_elements(sub_item, rep, axis))
-    return P.Concat(axis)(repeated_subs)
+        >>> from mindspore import Tensor
+        >>> a = tensor([0, 1, 2])
+        >>> output = a.repeat(2, 2)  # same as a.repeat((2, 2))
+        >>> print(output)
+        [[0 1 2 0 1 2]
+         [0 1 2 0 1 2]]
+    """
+    # only simulate 2 overload of repeat. Further check by F.tile
+    if repeats is None:
+        # no `repeats`: called by positional arguments like ``x.repeat(...)``
+        if len(args) == 1 and isinstance(args[0], (list, tuple)):
+            repeats = tuple(args[0])  # transform ``x.repeat([x0, x1, ...])`` (list type) to tuple
+        else:
+            repeats = args  # called as variable-length parameter like ``x.repeat(x0, x1, ...)``
+    else:
+        if args:  # simulate an exception thrown by Python interpreter
+            raise TypeError("repeat() got multiple values for argument 'repeat'")
+        # transform named argument with list type like ``x.repeat(repeats=[x0, x1, ...])`` to tuple
+        if isinstance(repeats, list):
+            repeats = tuple(repeats)
+    x_rank = F.rank(x)
+    if len(repeats) < x_rank:
+        raise ValueError(
+            "For repeat, number of items of repeats can not be smaller than the number of "
+            f"dimensions of self tensor, but got repeats with {len(repeats)}"
+            f" items and rank of self Tensor is {x_rank}.")
+    return F.tile(x, repeats)
 
 
 def repeat_interleave(x, repeats, dim=None):
@@ -2395,7 +2384,7 @@ def bool_func(*data):
 def cast_to_int(*data):
     target = data[0]
     if isinstance(target, (Tensor, Tensor_)):
-        target = Tensor(target, internal=True)
+        target = Tensor(target)
     if len(data) == 1:
         return int(target)
     return int(target, data[1])
@@ -2433,7 +2422,7 @@ def int_func(*data):
 @constexpr
 def cast_to_float(data):
     if isinstance(data, (Tensor, Tensor_)):
-        data = Tensor(data, internal=True)
+        data = Tensor(data)
     return float(data)
 
 
@@ -2496,6 +2485,18 @@ def tuple_func(data):
     return ret
 
 
+def dict_func(data):
+    """Implementation of `dict`."""
+    if isinstance(data, (tuple, list)):
+        keys = F.make_tuple()
+        values = F.make_tuple()
+        for pair in data:
+            keys = keys + F.make_tuple(pair[0])
+            values = values + F.make_tuple(pair[1])
+        return F.make_dict(keys, values)
+    raise TypeError('Currently, dict() only supports tuple or list input.')
+
+
 def ms_zip(*data):
     """Packs elements in the corresponding positions in multiple sequences into tuples."""
     x = ()
@@ -2605,7 +2606,7 @@ def ms_max_one_element(x):
 def ms_max(*data):
     """Implementation of `max`."""
     len_data = get_max_min_data_len(data)
-    if len_data <= 0: # pylint: disable=no-else-raise
+    if len_data <= 0:  # pylint: disable=no-else-raise
         raise TypeError("max() requires 1 argument at least.")
     elif len_data == 1:
         x = data[0]
@@ -2681,7 +2682,7 @@ def ms_min_one_element(x):
 def ms_min(*data):
     """Implementation of `min`."""
     len_data = get_max_min_data_len(data)
-    if len_data <= 0: # pylint: disable=no-else-raise
+    if len_data <= 0:  # pylint: disable=no-else-raise
         raise TypeError("min() requires 1 argument at least.")
     elif len_data == 1:
         x = data[0]
@@ -3214,7 +3215,7 @@ def random_categorical(x, num_sample, seed=0, dtype=mstype.int64):
 @constexpr
 def empty_tensor(dtype):
     """Return empty tensor"""
-    return Tensor_([], dtype)
+    return Tensor([], dtype)
 
 
 @constexpr
@@ -3291,7 +3292,7 @@ check_bool = constexpr(validator.check_bool)
 @constexpr
 def empty_compile(dtype, shape):
     """Returns an empty Tensor."""
-    return Tensor_(dtype, shape)
+    return Tensor(dtype=dtype, shape=shape)
 
 
 def tensor_bool(x):
@@ -3414,8 +3415,8 @@ def normal_(input, mean=0, std=1, *, generator=None):
     """
     if generator is None:
         generator = default_generator
-    seed, offset = generator._step(  # pylint: disable=protected-access
-        generator_step_)
+
+    seed, offset = generator._step(generator_step_)
     return inplace_normal_op(input, mean, std, seed, offset)
 
 
@@ -3698,11 +3699,11 @@ def sparse_ndim_(x):
     return F.tuple_len(x.shape)
 
 
-def bernoulli(input, p=0.5, seed=None):
+def bernoulli(input, *, generator=None):
     """
     Randomly draws binary numbers from a Bernoulli distribution.
     """
-    return F.bernoulli(input, p, seed)
+    return F.bernoulli_ext(input, generator=generator)
 
 
 def gather_nd(input_x, indices):
@@ -3935,7 +3936,7 @@ def atanh(x):
     return F.atanh(x)
 
 
-def baddbmm(x, batch1, batch2, beta=1, alpha=1):
+def baddbmm(x, batch1, batch2, *, beta=1, alpha=1):
     r"""
     For details, please refer to :func:`mindspore.ops.baddbmm`.
     """
@@ -4103,6 +4104,16 @@ def erfinv(input):
     return F.erfinv(input)
 
 
+def erfinv_(input):
+    r"""
+    For details, please refer to :func:`mindspore.Tensor.erfinv_`.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+    """
+    return inplace_erfinv_op(input)
+
+
 def less_equal(input, other):
     r"""
     Computes the boolean value of :math:`input\_x <= other` element-wise.
@@ -4467,6 +4478,29 @@ def uniform(input, from_=0., to=1., generator=None):
     return F.uniform_ext(input, from_, to, generator)
 
 
+def uniform_(input, from_=0, to=1, *, generator=None):
+    r"""
+    For details, please refer to :func:`mindspore.Tensor.uniform_`.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+    """
+    if generator is None:
+        generator = default_generator
+    seed, offset = generator._step(generator_step_) # pylint: disable=protected-access
+    return inplace_uniform_op(input, from_, to, seed, offset)
+
+
+def exponential_(input, lambd=1, *, generator=None):
+    r"""
+        Fills `self` tensor with elements drawn from the exponential distribution:
+    """
+    if generator is None:
+        generator = default_generator
+    seed, offset = generator._step(generator_step_) # pylint: disable=protected-access
+    return inplace_exponential_op(input, lambd, seed, offset)
+
+
 def amin(input, axis=None, keep_dims=False):
     r"""
     For details, please refer to :func:`mindspore.ops.amin`.

mindspore/_extends/pijit/__init__.py

@@ -18,6 +18,6 @@ Helper module for pijit analyze
 
 
 from .pijit_func_white_list import _func_map as pijit_func_white_list_map
+from .tensor_func_list import get_tensor_method_name
 
-
-__all__ = ['pijit_func_white_list_map']
+__all__ = ['pijit_func_white_list_map', "get_tensor_method_name"]

mindspore/_extends/pijit/pijit_func_white_list.py

@@ -29,7 +29,7 @@ from mindspore.common.api import jit
 from mindspore.common.tensor import Tensor
 from mindspore.common._register_for_tensor import Registry
 from mindspore._c_expression import MetaFuncGraph_, function_id
-from mindspore._c_expression import Tensor as Tensor_
+from mindspore._c_expression import TensorPy as Tensor_
 from mindspore._extends.parse.resources import convert_object_map
 from mindspore import _checkparam as validator
 from mindspore import Parameter, ParameterTuple
@@ -49,11 +49,12 @@ from mindspore.train.data_sink import _init_sink_dataset
 from mindspore.train.summary import SummaryRecord
 from mindspore.train._utils import _exec_datagraph
 from mindspore.train.summary.writer import BaseWriter
-from mindspore.train.serialization import _exec_save, load, export_split_mindir, obfuscate_model, _parse_ckpt_proto, \
+from mindspore.train.serialization import _exec_save, load, export_split_mindir, _parse_ckpt_proto, \
     _generate_front_info_for_param_data_file, _get_data_file, _encrypt_data, _split_save, _save_mindir_together, \
     _load_into_param_dict
 from mindspore.parallel import _cost_model_context
 from mindspore.parallel._offload_context import offload_context
+from mindspore.parallel._utils import _is_in_data_parallel_mode
 from mindspore.run_check._check_version import check_version_and_env_config
 from mindspore.dataset.callback.ds_callback import DSCallback, WaitedDSCallback
 from mindspore.dataset.transforms.c_transforms import TensorOperation as CTensorOperation, OneHot as COneHot, \
@@ -360,6 +361,7 @@ FUNC_KEY_DICT_ITEMS = 22  # dict.items
 FUNC_KEY_PRIMITIVE_ASSIGN = 23  # mindspore.ops.assign, Primitive("Assign")
 FUNC_KEY_TENSOR_SETITEM = 24  # Tensor.__setitem__
 FUNC_KEY_TENSOR_ASSIGN_VALUE = 25  # Tensor.assign_value
+FUNC_KEY_TENSOR_IS_CONTIGUOUS = 26  # Tensor.is_contiguous
 
 # Initialized only once. This map will initialize by c++ when start pijit.
 # key is customer if fuzzy match. (Primitive, constexpr, primexpr, MetaFuncGraph)
@@ -376,19 +378,19 @@ _func_map = {
     constexpr_key: FUNC_KEY_CONSTEXPR,
     primexpr_key: FUNC_KEY_PRIMEXPR,
     meta_func_graph_key: FUNC_KEY_META_FUNCG_RAPH,
-    id(GraphCell.__call__): FUNC_KEY_GRAPH_CELL,
+    function_id(GraphCell.__call__): FUNC_KEY_GRAPH_CELL,
     id(psjit_code): FUNC_KEY_PSJIT_CODE,
-    id(_get_cache_prim): FUNC_KEY_GET_CACHE_PRIM,
-    id(Registry.get): FUNC_KEY_REGISTRY_GET,
+    function_id(_get_cache_prim): FUNC_KEY_GET_CACHE_PRIM,
+    function_id(Registry.get): FUNC_KEY_REGISTRY_GET,
 
     # tensor side-effect
     primitive_assign_key: FUNC_KEY_PRIMITIVE_ASSIGN,
-    id(F.assign): FUNC_KEY_PRIMITIVE_ASSIGN,
-    id(Tensor.assign_value): FUNC_KEY_TENSOR_ASSIGN_VALUE,
-    id(Tensor.__setitem__): FUNC_KEY_TENSOR_SETITEM,
+    function_id(F.assign): FUNC_KEY_PRIMITIVE_ASSIGN,
+    function_id(Tensor.assign_value): FUNC_KEY_TENSOR_ASSIGN_VALUE,
+    function_id(Tensor.__setitem__): FUNC_KEY_TENSOR_SETITEM,
 
     # Tensor method
-    id(Tensor.astype): FUNC_KEY_TENSOR_ASTYPE,
+    function_id(Tensor.astype): FUNC_KEY_TENSOR_ASTYPE,
 
     # types.BuiltinFunctionType
     function_id(isinstance): FUNC_KEY_BUILTIN_FUNC,
@@ -448,6 +450,7 @@ _func_map = {
     function_id(str.isalnum): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.isidentifier): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.isprintable): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.replace): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.format): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.format_map): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.__format__): FUNC_KEY_BUILTIN_FUNC,
@@ -472,7 +475,7 @@ _func_map = {
     function_id(Tensor_.getitem_index_info): FUNC_KEY_BUILTIN_FUNC,
     function_id(Tensor_.get_bytes): FUNC_KEY_BUILTIN_FUNC,
     function_id(Tensor_.is_init): FUNC_KEY_BUILTIN_FUNC,
-    function_id(Tensor_.is_contiguous): FUNC_KEY_BUILTIN_FUNC,
+    function_id(Tensor_.is_contiguous): FUNC_KEY_TENSOR_IS_CONTIGUOUS,
     function_id(Tensor_.stride): FUNC_KEY_BUILTIN_FUNC,
     # Tensor_.asnumpy need real tensor value
 
@@ -488,6 +491,7 @@ _func_map = {
     function_id(validator.check_number_range): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(validator.check_is_int): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(validator.check_is_number): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(validator.check_positive_int_sequence): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(np_version_valid): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(_is_initialized): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(_set_elegant_exit_handle): FUNC_KEY_PIJIT_CONSTEXPR,
@@ -496,7 +500,9 @@ _func_map = {
     function_id(get_rank_size): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(get_rank_id): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(offload_context): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(_is_in_data_parallel_mode): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(check_version_and_env_config): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(Tensor.tolist): FUNC_KEY_PIJIT_CONSTEXPR,
 
     # inner function
     function_id(type_size_in_bytes): FUNC_KEY_BUILTIN_FUNC,
@@ -530,7 +536,6 @@ _func_map = {
     function_id(_exec_save): FUNC_KEY_PIJIT_FORBIDDEN,
     function_id(load): FUNC_KEY_PIJIT_FORBIDDEN,
     function_id(export_split_mindir): FUNC_KEY_PIJIT_FORBIDDEN,
-    function_id(obfuscate_model): FUNC_KEY_PIJIT_FORBIDDEN,
     function_id(_parse_ckpt_proto): FUNC_KEY_PIJIT_FORBIDDEN,
     function_id(_generate_front_info_for_param_data_file): FUNC_KEY_PIJIT_FORBIDDEN,
     function_id(_get_data_file): FUNC_KEY_PIJIT_FORBIDDEN,

mindspore/_extends/pijit/tensor_func_list.py

@@ -0,0 +1,27 @@
+# Copyright 2025 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.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Store and get tensor method"""
+from mindspore import Tensor
+from mindspore._c_expression import function_id
+
+tensor_method_id_to_name = {}
+for method_name in dir(Tensor):
+    method_id = function_id(getattr(Tensor, method_name))
+    tensor_method_id_to_name[method_id] = method_name
+
+
+def get_tensor_method_name(id):
+    """Get method name by function id"""
+    return tensor_method_id_to_name.get(id, None)

mindspore/_extends/utils.py

@@ -28,7 +28,7 @@ def cell_attr_register(fn=None, attrs=None):
 
     Args:
         fn (function): __init__ function of cell.
-        attrs (list(string) | string): attr list.
+        attrs (list(str) | str): attr list.
 
     Returns:
         function, original function.

mindspore/amp.py

@@ -99,12 +99,12 @@ def _grad_scale(scale, grad):
     return grad * scale.astype(grad.dtype)
 
 
-@jit
+@jit(backend="ms_backend")
 def _grad_scale_map(scale_value, inputs):
     return _hypermap(_partial(_grad_scale, scale_value), inputs)
 
 
-@jit
+@jit(backend="ms_backend")
 def _grad_unscale_map(scale_value, inputs):
     return _hypermap(_partial(_grad_unscale, scale_value), inputs)
 
@@ -116,7 +116,7 @@ def _overflow(inputs):
     return 1 - status.all()
 
 
-@jit
+@jit(backend="ms_backend")
 def _all_finite(inputs, check_overflow_mode, enable_allfinite):
     """all finite check"""
     if _ascend_target():
@@ -325,7 +325,7 @@ class StaticLossScaler(LossScaler):
 
 class DynamicLossScaler(LossScaler):
     r"""
-    Dynamic Loss scale class.
+    Manager for dynamically adjusting the loss scaling factor.
 
     Dynamic loss scaling tries to determine the largest loss scale value that
     will keep gradients finite. It does this by increasing the loss scale every

mindspore/boost/__init__.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 # ============================================================================
 """
-Boost provide auto accelerating for network, such as Less BN, Gradient Freeze, Gradient
-accumulation and so on.
+Boost is able to automatically optimize network performance, e.g., by reducing BN, gradient freezing,
+and accumulating gradients to achieve network acceleration.
 
 Note:
     This feature is a beta feature, and we are still improving its functionality.

mindspore/boost/base.py

@@ -21,15 +21,12 @@ import math
 import copy
 import numpy as np
 from scipy import linalg as la
-from mindspore.context import ParallelMode
 import mindspore.nn as nn
 from mindspore.nn.optim import LARS
 from mindspore import log as logger
 from mindspore.common import Parameter
-from mindspore.communication.management import get_group_size
+from mindspore.communication.management import get_rank, get_group_size
 from mindspore.train.serialization import load_checkpoint
-from mindspore.parallel._utils import _get_global_rank
-from mindspore.parallel._auto_parallel_context import auto_parallel_context
 from mindspore.boost.less_batch_normalization import CommonHeadLastFN
 
 
@@ -329,7 +326,7 @@ def _get_local_pca_mat_path(weight_load_dir, pca_mat_path, n_component, device_n
     if os.path.exists(save_pca_end_path):
         os.remove(save_pca_end_path)
 
-    rank = _get_global_rank()
+    rank = get_rank()
     local_pca_mat_path = full_pca_mat_path[:-4] + "_rank_" + str(rank) + ".npy"
     if os.path.exists(local_pca_mat_path):
         os.remove(local_pca_mat_path)
@@ -498,8 +495,7 @@ def _save_local_pca_mat(pca_mat, full_pca_mat_path, n_component):
         full_pca_mat_path (str): the path of full pca mat.
         n_component (int): pca component.
     """
-    parallel_mode = auto_parallel_context().get_parallel_mode()
-    rank_size = 1 if parallel_mode == ParallelMode.STAND_ALONE else get_group_size()
+    rank_size = get_group_size()
     local_dim = math.ceil(n_component // rank_size)
     for rank_id in range(rank_size):
         start_index = rank_id * local_dim

mindspore/boost/boost_cell_wrapper.py

@@ -21,7 +21,7 @@ from mindspore.nn.wrap import TrainOneStepCell
 import mindspore.context as context
 from mindspore.context import ParallelMode
 from mindspore.parallel._utils import _get_global_rank, _get_device_num, _get_gradients_mean
-from mindspore.communication.management import get_group_size, create_group
+from mindspore.communication.management import get_rank, get_group_size, create_group
 from mindspore.nn.cell import Cell
 from mindspore.nn import SequentialCell
 from mindspore.common import Tensor
@@ -388,7 +388,7 @@ class BoostTrainOneStepCell(TrainOneStepCell):
         gamma = self.auto_boost.gamma
         alpha = self.auto_boost.alpha
         sigma = self.auto_boost.sigma
-        _rank = _get_global_rank()
+        _rank = get_rank()
         _rank_size = 1 if self.parallel_mode == ParallelMode.STAND_ALONE else get_group_size()
         n_components = self.auto_boost.n_components
         timeout = self.auto_boost.timeout