mindspore 2.4.0__cp310-none-any.whl → 2.4.10__cp310-none-any.whl

mindspore/mint/__init__.py

@@ -15,6 +15,7 @@
 """mint module."""
 from __future__ import absolute_import
 import mindspore.ops as ops
+from mindspore.ops.primitive import constexpr
 from mindspore.common._register_for_tensor import tensor_operator_registry_for_mint
 from mindspore.common.tensor import Tensor
 from mindspore.ops.function.array_func import gather_ext as gather, max_ext as max, min_ext as min
@@ -320,8 +321,13 @@ from mindspore.ops.auto_generate import erfc
 from mindspore.ops.auto_generate import expm1
 # 208
 from mindspore.ops.function.array_func import eye
+from mindspore.ops.function.random_func import randperm_ext as randperm
 from mindspore.ops.function.random_func import rand_ext as rand
 from mindspore.ops.function.random_func import rand_like_ext as rand_like
+from mindspore.ops.function.random_func import randn_ext as randn
+from mindspore.ops.function.random_func import randn_like_ext as randn_like
+from mindspore.ops.function.random_func import randint_ext as randint
+from mindspore.ops.function.random_func import randint_like_ext as randint_like
 # 210
 from mindspore.ops.auto_generate import floor
 # 231
@@ -364,6 +370,9 @@ from mindspore.ops.auto_generate import logaddexp_ext as logaddexp
 # 610
 from mindspore.ops.function.math_func import nan_to_num
 
+# 695
+from mindspore.ops.auto_generate import count_nonzero
+
 
 def add(input, other, *, alpha=1):
     r"""
@@ -661,6 +670,474 @@ def cummax(input, dim):
     return ops.auto_generate.cummax(input, dim)
 
 
+def _einsum_convert_sublist_to_label(num, ell_num=False):
+    """Convert sublist to label."""
+    if num == Ellipsis or ell_num and num == 52:
+        return '...'
+    if 0 <= num < 26:
+        return chr(num + ord('A'))
+    if 26 <= num < 52:
+        return chr(num + ord('a') - 26)
+    raise ValueError(f'For einsum, the number in sublist must be in range [0, 52), but got {num}')
+
+
+def _einsum_convert_label_to_index(label):
+    """Convert label to index."""
+    label_num = ord(label)
+    if ord('A') <= label_num <= ord('Z'):
+        return label_num - ord('A')
+    if ord('a') <= label_num <= ord('z'):
+        return label_num - ord('a') + 26
+    if label_num == ord('.'):
+        return 52
+    raise ValueError(f'For einsum, the label in equation must be in [a-zA-Z] or ., but got {label}')
+
+
+def _einsum_convert_sublist(equation, *operands):
+    """Convert the sublist to an equation operand if the received input is a sublist format."""
+    if isinstance(equation, Tensor):
+        equation_tmp = ''
+        for i, lst in enumerate(operands):
+            if i % 2 == 0:
+                for _, num in enumerate(lst):
+                    equation_tmp += _einsum_convert_sublist_to_label(num)
+                if i in (len(operands) - 1, len(operands) - 2):
+                    continue
+                equation_tmp += ','
+        if len(operands) % 2 == 0:
+            equation_tmp += '->'
+            for _, num in enumerate(operands[-1]):
+                equation_tmp += _einsum_convert_sublist_to_label(num)
+            operands_tmp = list([equation]) + list(operands[1:-1:2])
+        else:
+            operands_tmp = list([equation]) + list(operands[1::2])
+        equation = equation_tmp
+        operands = tuple(operands_tmp)
+    if len(operands) == 0: # pylint: disable=len-as-condition
+        raise ValueError("For einsum, the 'operands' must have at least one operand.")
+    return equation, operands
+
+
+def _einsum_check_inputargs(equation, operands):
+    """Check equation and operands."""
+    if not isinstance(equation, str):
+        raise TypeError(f"For einsum, 'equation' must be a str, but got {type(equation)}.")
+    for operand in operands:
+        if not isinstance(operand, Tensor):
+            raise TypeError(f"For einsum, members of 'operands' must be Tensor, but got {type(operand)}.")
+
+
+@constexpr
+def _einsum_parse_equation(equation):
+    """Parse equation."""
+    l_equation = ''
+    r_equation = ''
+    equation = equation.replace(' ', '')
+
+    if '->' in equation:
+        l_equation, r_equation = equation.split('->', 1)
+        if l_equation == '':
+            raise ValueError('For einsum, equation must contain characters to the left fo the arrow.')
+    else:
+        l_equation = equation
+
+    if ',' in l_equation:
+        l_equationlst = l_equation.split(",")
+    else:
+        l_equationlst = [l_equation]
+
+    l_equationlst = []
+
+    for subequation in l_equation.split(','):
+        if '.' in subequation and ('...' not in subequation or subequation.count('.') != 3):
+            raise ValueError(f"For einsum, an ellipsis in the equation must include three continuous \'.\', "
+                             f"and can only be found once.")
+        subequation_lst = [_einsum_convert_label_to_index(label) for label in subequation.replace('...', '.')]
+        l_equationlst.append(subequation_lst)
+
+    if "." in r_equation and ('...' not in r_equation or r_equation.count('.') != 3):
+        raise ValueError(f"For einsum, an ellipsis in the equation must include three continuous \'.\', "
+                         f"and can only be found once.")
+    r_equationlst = [_einsum_convert_label_to_index(label) for label in r_equation.replace('...', '.')]
+
+    return l_equationlst, r_equationlst, ('->' in equation)
+
+
+def _einsum_parse_labels(l_equationlst, operands):
+    """Parse left script of equation."""
+    align_rank = 0
+    max_labels = 53
+    labels_count = [0] * max_labels
+    labels2dimlst = [None] * max_labels
+
+    if len(operands) != len(l_equationlst):
+        raise ValueError(f"For einsum, 'operands' is not equal to specified in the 'equation', "
+                         f"but got {len(operands)} and {len(l_equationlst)}.")
+
+    for idx, sub_equ in enumerate(l_equationlst):
+        start_dim = 0
+        label_num = 0
+        operand_shape = list(operands[idx].shape)
+        for label in sub_equ:
+            label_num += 1
+            end_dim = start_dim + 1
+
+            # Label is ellipsis
+            if label == 52:
+                end_dim = len(operand_shape) - len(sub_equ) + label_num
+            if labels2dimlst[label] is None:
+                labels2dimlst[label] = operand_shape[start_dim:end_dim]
+                align_rank += (end_dim - start_dim)
+            else:
+                if labels2dimlst[label] != operand_shape[start_dim:end_dim]:
+                    raise ValueError(f"For einsum, one label in 'equation' can only represent the same dimension "
+                                     f"in 'operands', but '{_einsum_convert_sublist_to_label(label, True)}' "
+                                     f"represented different dimensions.")
+            labels_count[label] += 1
+            start_dim = end_dim
+        if label_num != len(sub_equ) or start_dim != len(operand_shape):
+            raise ValueError(f"For einsum, the numbers of labels specified in the 'equation' does not match "
+                             f"'operands[{idx}]'.")
+    return labels2dimlst, labels_count, align_rank
+
+
+def _einsum_infer_output(r_equationlst, arrow_exist, labels2dimlst, labels_count):
+    """Parse right script of equation and infer output shape."""
+    idx = 0
+    idle_idx = -1
+    output_shape = []
+    labels_perm_idx = [idle_idx] * 53
+
+    if arrow_exist:
+        for label in r_equationlst:
+            if labels_count[label] != 0:
+                output_shape += labels2dimlst[label]
+                if labels_perm_idx[label] != idle_idx:
+                    raise ValueError(f"For einsum, '{_einsum_convert_sublist_to_label(label, True)}' or {label} in "
+                                     f"sublist format has appears more than once in output subscript.")
+                labels_perm_idx[label] = idx
+                idx += len(labels2dimlst[label])
+            else:
+                raise ValueError(f"For einsum, the label to the right of arrow in the 'equation' must appear on "
+                                 f"left, but '{_einsum_convert_sublist_to_label(label, True)}' does not.")
+    else:
+        if labels_count[52] != 0:
+            output_shape += labels2dimlst[52]
+            labels_perm_idx[52] = idx
+            idx += len(labels2dimlst[52])
+        for label, count in enumerate(labels_count):
+            if count == 1:
+                output_shape += labels2dimlst[label]
+                labels_perm_idx[label] = idx
+                idx += len(labels2dimlst[label])
+
+    for label, count in enumerate(labels_count):
+        if count != 0 and labels_perm_idx[label] == idle_idx:
+            labels_perm_idx[label] = idx
+            idx += 1
+
+    return output_shape, labels_perm_idx
+
+
+def _einsum_adjust_operands(operands, l_equationlst, labels2dimlst, labels_perm_idx, align_rank):
+    """Align operands to output as possible."""
+    # Unsqueeze miss dimensions to make all operands has same rank, compute diagonal if operand has same label.
+    # Then use _labels_perm_idx to transpose all operands to align dimensions with output.
+    adjust_operands = []
+    for idx, operand in enumerate(operands):
+        idle_dim = -1
+        align_axis = [idle_dim] * align_rank
+        label_dims = [idle_dim] * 53
+        dim = 0
+
+        for label in l_equationlst[idx]:
+            if label_dims[label] != idle_dim:
+                operand = ops.diagonal(operand, 0, label_dims[label], dim)
+                diag_perm = []
+                diag_dim = 0
+                for i in range(len(operand.shape)):
+                    if i == label_dims[label]:
+                        diag_perm.append(len(operand.shape) - 1)
+                    else:
+                        diag_perm.append(diag_dim)
+                        diag_dim += 1
+                operand = permute(operand, tuple(diag_perm))
+            else:
+                label_dims[label] = dim
+                if label == 52:
+                    for ell_idx in range(len(labels2dimlst[label])):
+                        align_axis[labels_perm_idx[label] + ell_idx] = dim
+                        dim += 1
+                else:
+                    align_axis[labels_perm_idx[label]] = dim
+                    dim += 1
+        if len(operand.shape) < align_rank:
+            for i, axis in enumerate(align_axis):
+                if axis == idle_dim:
+                    align_axis[i] = dim
+                    dim += 1
+            missing_dims = [1] * (align_rank - len(operand.shape))
+            operand_shape = list(operand.shape) + missing_dims
+            operand = reshape(operand, operand_shape)
+        operand = permute(operand, tuple(align_axis))
+        adjust_operands.append(operand)
+    return adjust_operands
+
+
+def _einsum_find_dimlastop(align_rank, operands, adjust_operands):
+    """Find dim last operand."""
+    dim_last_op = [0 for _ in range(align_rank)]
+    has_zero_dim = False
+    for dim in range(align_rank):
+        broadcast_dim = adjust_operands[0].shape[dim]
+        for idx in range(1, len(adjust_operands)):
+            other_dim = adjust_operands[idx].shape[dim]
+            if broadcast_dim != other_dim and broadcast_dim != 1 and other_dim != 1:
+                err_msg = "For einsum, operands do not broadcast after align to output [shapes :origin -> adjust]:"
+                for i in range(len(operands)):
+                    err_msg += f" {operands[i].shape} -> {adjust_operands[i].shape}"
+                raise ValueError(err_msg)
+            if other_dim != 1:
+                dim_last_op[dim] = idx
+                broadcast_dim = other_dim
+        has_zero_dim = has_zero_dim or broadcast_dim == 0
+    return dim_last_op, has_zero_dim
+
+
+def _einsum_multiplication(sum_dims, l_tensor, r_tensor):
+    """Compute bmm for einsum."""
+    batch_dims = []
+    lonly_dims = []
+    ronly_dims = []
+    batch_size = 1
+    lonly_size = 1
+    ronly_size = 1
+    sum_size = 1
+
+    l_shape = l_tensor.shape
+    r_shape = r_tensor.shape
+
+    # Compute sum if dim is in sum_dims and get shapes for bmm
+    for i in range(len(l_shape)):
+        sum_l = l_shape[i] > 1
+        sum_r = r_shape[i] > 1
+        if i in sum_dims:
+            if sum_l and sum_r:
+                sum_size *= l_shape[i]
+            elif sum_l:
+                l_tensor = sum(l_tensor, i, True)
+            elif sum_r:
+                r_tensor = sum(r_tensor, i, True)
+        elif sum_l and sum_r:
+            batch_dims.append(i)
+            batch_size *= l_shape[i]
+        elif sum_l:
+            lonly_dims.append(i)
+            lonly_size *= l_shape[i]
+        else:
+            ronly_dims.append(i)
+            ronly_size *= r_shape[i]
+
+    # Compute the einsum bmm operators pipeline.
+    # The whole operators pipline is transpose(in) -> reshape(in) -> bmm(in) -> reshape(out) -> transpose(out).
+    l_reshape_shape = (batch_size, lonly_size, sum_size)
+    r_reshape_shape = (batch_size, sum_size, ronly_size)
+
+    out_reshape_shape = [l_shape[dim] for dim in batch_dims]
+    out_reshape_shape += [l_shape[dim] for dim in lonly_dims]
+    out_reshape_shape += [1 for _ in sum_dims]
+    out_reshape_shape += [r_shape[dim] for dim in ronly_dims]
+
+    l_perm_axis = batch_dims + lonly_dims + sum_dims + ronly_dims
+    r_perm_axis = batch_dims + sum_dims + ronly_dims + lonly_dims
+    out_perm_axis = [-1] * len(out_reshape_shape)
+
+    out_dim = 0
+    for idx in range(len(l_perm_axis)):
+        out_perm_axis[l_perm_axis[idx]] = out_dim
+        out_dim += 1
+
+    l_tensor = permute(l_tensor, tuple(l_perm_axis))
+    l_tensor = reshape(l_tensor, l_reshape_shape)
+
+    r_tensor = permute(r_tensor, tuple(r_perm_axis))
+    r_tensor = reshape(r_tensor, r_reshape_shape)
+
+    output = bmm(l_tensor, r_tensor)
+    output = reshape(output, out_reshape_shape)
+    output = permute(output, tuple(out_perm_axis))
+
+    output_origin_shape = output.shape
+    output_squeeze_shape = []
+    for dim in range(len(output_origin_shape)):
+        if dim not in sum_dims:
+            output_squeeze_shape.append(output_origin_shape[dim])
+
+    return reshape(output, output_squeeze_shape)
+
+
+def _einsum_squeeze(operand, dim):
+    '''Will be replaced by mint.squeeze in the future'''
+    operand_shape = operand.shape
+    squeeze_shape = []
+    for idx in range(len(operand_shape)):
+        if idx != dim:
+            squeeze_shape.append(operand_shape[idx])
+    return reshape(operand, squeeze_shape)
+
+
+def _einsum(equation, operands):
+    '''Einsum main process'''
+    _l_equationlst, _r_equationlst, _arrow_exist = _einsum_parse_equation(equation)
+    _labels2dimlst, _labels_count, _align_rank = _einsum_parse_labels(_l_equationlst, operands)
+    _output_shape, _labels_perm_idx = _einsum_infer_output(_r_equationlst, _arrow_exist, _labels2dimlst, _labels_count)
+    _output_rank = len(_output_shape)
+
+    _adjust_operands = _einsum_adjust_operands(operands, _l_equationlst, _labels2dimlst, _labels_perm_idx, _align_rank)
+    _dim_last_op, _has_zero_dim = _einsum_find_dimlastop(_align_rank, operands, _adjust_operands)
+    _result = _adjust_operands[0]
+
+    # Fast path if operands has zero dim.
+    if _has_zero_dim:
+        return zeros(_output_shape, dtype=_result.dtype)
+
+    # Sum or squeeze dimensions that is 1 for all rest operands.
+    _reduce_dim = _output_rank
+    for dim in range(_output_rank, _align_rank):
+        if _dim_last_op[dim] == 0:
+            if _result.shape[_reduce_dim] == 1:
+                _result = _einsum_squeeze(_result, _reduce_dim)
+            else:
+                _result = sum(_result, _reduce_dim)
+        else:
+            _reduce_dim += 1
+
+    # Compute multiplication if operands are more than two.
+    for i in range(1, len(_adjust_operands)):
+        operand = _adjust_operands[i]
+        dim = _output_rank
+        sum_dims = []
+        for j in range(_output_rank, _align_rank):
+            if _dim_last_op[j] < i:
+                operand = _einsum_squeeze(operand, dim)
+            elif _dim_last_op[j] == i:
+                if _result.shape[dim] == 1:
+                    operand = sum(operand, dim)
+                    _result = _einsum_squeeze(_result, dim)
+                else:
+                    sum_dims.append(dim)
+                    dim += 1
+            else:
+                dim += 1
+
+        if sum_dims == []:
+            _result = mul(_result, operand)
+        elif len(sum_dims) == len(_result.shape):
+            _result = ops.auto_generate.dot(flatten(_result), flatten(operand))
+        else:
+            _result = _einsum_multiplication(sum_dims, _result, operand)
+
+    return _result
+
+
+def einsum(equation, *operands):
+    r"""
+    According to the Einstein summation Convention (Einsum),
+    the product of the input tensor elements is summed along the specified dimension.
+    You can use this operator to perform diagonal, reducesum, transpose, matmul, mul, inner product operations, etc.
+
+    Note:
+        The sublist format is also supported. For example, mint.einsum(op1, sublist1, op2, sublist2, ..., sublist_out).
+        In this format, equation can be derived by the sublists which are made up of Python's Ellipsis and list of
+        integers in [0, 52). Each operand is followed by a sublist and an output sublist is at the end.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        equation (str): Notation based on the Einstein summation convention, represent the operation you want to do.
+            the value can contain only letters, commas, ellipsis and arrow.
+            The letters represent input tensor dimension, commas represent separate tensors, ellipsis indicates
+            the tensor dimension that you do not care about, the left of the arrow indicates the input tensors,
+            and the right of it indicates the desired output dimension.
+        operands (Tensor): Input tensor used for calculation. The dtype of the tensor must be the same.
+
+    Returns:
+        Tensor, the shape of it can be obtained from the `equation` , and the dtype is the same as input tensors.
+
+    Raises:
+        TypeError: If `equation` is invalid, or the `equation` does not match the input tensor.
+        ValueError: If the number in sublist is not in [0, 52) in sublist format.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor, mint
+        >>> x = Tensor(np.array([1.0, 2.0, 4.0]), mindspore.float32)
+        >>> equation = "i->"
+        >>> output = mint.einsum(equation, x)
+        >>> print(output)
+        [7.]
+        >>> x = Tensor(np.array([1.0, 2.0, 4.0]), mindspore.float32)
+        >>> y = Tensor(np.array([2.0, 4.0, 3.0]), mindspore.float32)
+        >>> equation = "i,i->i"
+        >>> output = mint.einsum(equation, x, y)
+        >>> print(output)
+        [ 2. 8. 12.]
+        >>> x = Tensor(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), mindspore.float32)
+        >>> y = Tensor(np.array([[2.0, 3.0], [1.0, 2.0], [4.0, 5.0]]), mindspore.float32)
+        >>> equation = "ij,jk->ik"
+        >>> output = mint.einsum(equation, x, y)
+        >>> print(output)
+        [[16. 22.]
+         [37. 52.]]
+        >>> x = Tensor(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), mindspore.float32)
+        >>> equation = "ij->ji"
+        >>> output = mint.einsum(equation, x)
+        >>> print(output)
+        [[1. 4.]
+         [2. 5.]
+         [3. 6.]]
+        >>> x = Tensor(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), mindspore.float32)
+        >>> equation = "ij->j"
+        >>> output = mint.einsum(equation, x)
+        >>> print(output)
+        [5. 7. 9.]
+        >>> x = Tensor(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), mindspore.float32)
+        >>> equation = "...->"
+        >>> output = mint.einsum(equation, x)
+        >>> print(output)
+        [21.]
+        >>> x = Tensor(np.array([1.0, 2.0, 3.0]), mindspore.float32)
+        >>> y = Tensor(np.array([2.0, 4.0, 1.0]), mindspore.float32)
+        >>> equation = "j,i->ji"
+        >>> output = mint.einsum(equation, x, y)
+        >>> print(output)
+        [[ 2. 4. 1.]
+         [ 4. 8. 2.]
+         [ 6. 12. 3.]]
+        >>> x = mindspore.Tensor([1, 2, 3, 4], mindspore.float32)
+        >>> y = mindspore.Tensor([1, 2], mindspore.float32)
+        >>> output = mint.einsum(x, [..., 1], y, [..., 2], [..., 1, 2])
+        >>> print(output)
+        [[1. 2.]
+         [2. 4.]
+         [3. 6.]
+         [4. 8.]]
+    """
+    _equation, _operands = _einsum_convert_sublist(equation, *operands)
+    _einsum_check_inputargs(_equation, _operands)
+
+    for operand in _operands:
+        if ops.is_sequence_shape_unknown(operand.shape) or ops.is_sequence_value_unknown(operand.shape):
+            raise ValueError(f"For einsum, the element of 'operands' can't be dynamic shape or dynamic rank.")
+
+    return _einsum(_equation, _operands)
+
+
 def item(input):
     r"""
     Returns the value of this tensor as a standard Python number.
@@ -694,7 +1171,8 @@ def item(input):
     if not isinstance(input, Tensor):
         raise TypeError(f"the input must be a Tensor, but got {type(input)}")
     if input.size != 1:
-        raise RuntimeError("a Tensor with {} elements cannot be converted to Scalar".format(input.size))
+        raise RuntimeError(
+            "a Tensor with {} elements cannot be converted to Scalar".format(input.size))
     return input.asnumpy().item()
 
 
@@ -1283,6 +1761,7 @@ __all__ = [
     # 31
     'cummax',
     'cummin',
+    'einsum',
     'sub',
     # 33
     'split',
@@ -1518,8 +1997,13 @@ __all__ = [
 
     # 256
     'median',
+    'randperm',
     'rand',
     'rand_like',
+    'randn',
+    'randn_like',
+    'randint',
+    'randint_like',
     # 210
     'floor',
     # 231
@@ -1554,6 +2038,9 @@ __all__ = [
 
     # 610
     'nan_to_num',
+
+    # 695
+    'count_nonzero',
 ]
 
 setattr(tensor_operator_registry_for_mint, 'add', add)
@@ -1568,7 +2055,8 @@ setattr(tensor_operator_registry_for_mint, 'item', item)
 setattr(tensor_operator_registry_for_mint, 'max', max)
 setattr(tensor_operator_registry_for_mint, 'mean', mean)
 setattr(tensor_operator_registry_for_mint, 'min', min)
-setattr(tensor_operator_registry_for_mint, 'repeat_interleave', repeat_interleave)
+setattr(tensor_operator_registry_for_mint,
+        'repeat_interleave', repeat_interleave)
 setattr(tensor_operator_registry_for_mint, 'ne', ne)
 setattr(tensor_operator_registry_for_mint, 'round', round)
 setattr(tensor_operator_registry_for_mint, 'sin', sin)

mindspore/mint/nn/__init__.py

@@ -28,7 +28,7 @@ from mindspore.nn import EmbeddingExt as Embedding, MaxPool2dExt as MaxPool2d, L
 # 2
 
 # 3
-
+from mindspore.nn.layer.basic import Identity
 # 4
 
 # 5
@@ -529,7 +529,7 @@ __all__ = [
     # 2
 
     # 3
-
+    'Identity',
     # 4
 
     # 5

mindspore/mint/optim/adamw.py

@@ -22,32 +22,25 @@ from mindspore.common import dtype as mstype
 from mindspore.ops import auto_generate as gen
 from mindspore.experimental.optim.optimizer import Optimizer
 from mindspore import _checkparam as validator
-from mindspore import mint
 
-_optim_adamw_opt = C.MultitypeFuncGraph("optim_adamw_opt")
 hyper_map = C.HyperMap()
 
 
-@_optim_adamw_opt.register("Function", "Float", "Float", "Float", "Float", "Float", "Tensor", "Bool", "Bool", "Tensor",
-                           "Tensor", "Tensor", "Tensor", "Tensor")
 def _run_optim_adamw_amsgrad_opt(opt, beta1, beta2, lr, eps, weight_decay, step, amsgrad, maximize, parameters, grads,
                                  exp_avg, exp_avg_sq, max_exp_avg_sq):
     """Apply adamw optimizer to the weight parameter."""
     success = True
-    opt(parameters, exp_avg, exp_avg_sq, max_exp_avg_sq, P.Cast()(grads, F.dtype(parameters)), step, lr, beta1, beta2,
-        weight_decay, eps, amsgrad, maximize)
+    opt(parameters, exp_avg, exp_avg_sq, max_exp_avg_sq, grads, step, lr, beta1, beta2, weight_decay, eps, amsgrad,
+        maximize)
     return success
 
 
-@_optim_adamw_opt.register("Function", "Float", "Float", "Float", "Float", "Float", "Tensor", "Bool", "Bool", "Tensor",
-                           "Tensor", "Tensor", "Tensor")
 def _run_optim_adamw_opt(opt, beta1, beta2, lr, eps, weight_decay, step, amsgrad, maximize, parameters, grads, exp_avg,
                          exp_avg_sq):
     """Apply adamw optimizer to the weight parameter."""
     success = True
-    max_exp_avg_sq = mint.zeros_like(exp_avg)
-    opt(parameters, exp_avg, exp_avg_sq, max_exp_avg_sq, P.Cast()(grads, F.dtype(parameters)), step, lr, beta1, beta2,
-        weight_decay, eps, amsgrad, maximize)
+    opt(parameters, exp_avg, exp_avg_sq, exp_avg_sq, grads, step, lr, beta1, beta2, weight_decay, eps, amsgrad,
+        maximize)
     return success
 
 
@@ -177,7 +170,6 @@ class AdamW(Optimizer):
         self.state_step = Parameter(Tensor([-1], mstype.float32), "state_step")
         self.increase_tensor = Tensor(1, mstype.float32)
         self.assignadd = P.AssignAdd()
-        self.op_cast = P.Cast()
         self.adamw_opt = gen.AdamW()
 
     def construct(self, gradients):
@@ -191,13 +183,13 @@ class AdamW(Optimizer):
             grads = tuple(gradients[start_id: end_id])
 
             if group.get("amsgrad"):
-                self.hyper_map(F.partial(_optim_adamw_opt, self.adamw_opt, beta1, beta2, float(lr),
+                self.hyper_map(F.partial(_run_optim_adamw_amsgrad_opt, self.adamw_opt, beta1, beta2, float(lr),
                                          group.get("eps"), group.get("weight_decay"), self.state_step,
                                          group.get("amsgrad"), maximize),
                                self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],
                                self.exp_avg_sq[start_id: end_id], group.get("max_exp_avg_sq"))
             else:
-                self.hyper_map(F.partial(_optim_adamw_opt, self.adamw_opt, beta1, beta2, float(lr),
+                self.hyper_map(F.partial(_run_optim_adamw_opt, self.adamw_opt, beta1, beta2, float(lr),
                                          group.get("eps"), group.get("weight_decay"), self.state_step,
                                          group.get("amsgrad"), maximize),
                                self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],

mindspore/nn/__init__.py

@@ -31,6 +31,7 @@ from mindspore.nn.wrap import *
 from mindspore.nn.grad import Jvp, Vjp
 from mindspore.nn.sparse import *
 from mindspore.nn.reinforcement import *
+from mindspore.nn.utils import *
 
 __all__ = ["Cell", "GraphCell"]
 __all__.extend(layer.__all__)
@@ -43,5 +44,6 @@ __all__.extend(sparse.__all__)
 __all__.extend(learning_rate_schedule.__all__)
 __all__.extend(dynamic_lr.__all__)
 __all__.extend(reinforcement.__all__)
+__all__.extend(utils.__all__)
 
 __all__.sort()