mindspore 1.10.0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/nn/optim/adamax.py

@@ -17,16 +17,15 @@ from __future__ import absolute_import
 
 from mindspore.common import dtype as mstype
 from mindspore.common.initializer import initializer
-from mindspore.common.api import ms_function
+from mindspore.common.api import jit
 from mindspore.ops import operations as P
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
-from mindspore._checkparam import Rel
 
 _ada_max_opt = C.MultitypeFuncGraph("ada_max_opt")
 
@@ -44,8 +43,8 @@ def _check_param_value(beta1, beta2, eps, prim_name):
     validator.check_value_type("beta1", beta1, [float], prim_name)
     validator.check_value_type("beta2", beta2, [float], prim_name)
     validator.check_value_type("eps", eps, [float], prim_name)
-    validator.check_float_range(beta1, 0.0, 1.0, Rel.INC_NEITHER, "beta1", prim_name)
-    validator.check_float_range(beta2, 0.0, 1.0, Rel.INC_NEITHER, "beta2", prim_name)
+    validator.check_float_range(beta1, 0.0, 1.0, validator.INC_NEITHER, "beta1", prim_name)
+    validator.check_float_range(beta2, 0.0, 1.0, validator.INC_NEITHER, "beta2", prim_name)
     validator.check_positive_float(eps, "eps", prim_name)
 
 
@@ -196,7 +195,7 @@ class AdaMax(Optimizer):
 
         self.opt = P.ApplyAdaMax()
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         gradients = self.flatten_gradients(gradients)
         gradients = self.decay_weight(gradients)

mindspore/nn/optim/adasum.py

@@ -22,7 +22,7 @@ import math
 import mindspore.nn as nn
 import mindspore.log as logger
 from mindspore import context
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.nn.cell import Cell
 from mindspore.common.parameter import ParameterTuple, Parameter
 from mindspore.parallel._utils import _get_global_rank, _get_stage_device_num
@@ -47,7 +47,8 @@ def _update_parameters_adasum(delta_weight, update_delta_weight, parameter, old_
     shape = F.shape(delta_weight)
     update_delta_weight = reshape(update_delta_weight, shape)
     new_parameter = old_parameter - update_delta_weight
-    return P.Assign()(parameter, new_parameter)
+    P.Assign()(parameter, new_parameter)
+    return parameter
 
 
 @_reshape_grads.register("Tensor", "Tensor", "Function")
@@ -113,7 +114,7 @@ def _adasum_opt_forward_process(left_send, allreduce, parameter_divisibility, al
     if parameter_divisibility:
         delta_w = P.Squeeze()(delta_w)
         ori_len = F.shape(delta_w)[0]
-        divide_len = ori_len / 2
+        divide_len = ori_len // 2
         left_part = delta_w[:divide_len]
         right_part = delta_w[divide_len:]
     else:
@@ -377,7 +378,8 @@ def _get_delta_weight_process(new_parameter, old_parameter):
 
 @_save_weight.register("Tensor", "Tensor")
 def _save_weight_process(new_parameter, old_parameter):
-    return P.Assign()(new_parameter, old_parameter)
+    P.Assign()(new_parameter, old_parameter)
+    return new_parameter
 
 
 @_clone_weight.register("Tensor", "Tensor")
@@ -410,8 +412,8 @@ class AdaSumByGradWrapCell(Cell):
     .. math::
         \begin{array}{ll}
           w_{t+1}=w_{t} - \alpha \cdot Adasum(g_{1}, g_{2})  \\
-          w_{t+1}=w_{t} - \alpha \cdot [(1 - \frac{g_2^{T}\cdot g_1}{2\cdot \left \| g_1 \right \|^2 })\cdot g_1 +
-           (1 - \frac{g_1^{T}\cdot g_2}{2\cdot \left \| g_2 \right \|^2 })\cdot g_2]  \\
+          w_{t+1}=w_{t} - \alpha \cdot [(1 - \frac{g_2^{T}\cdot g_1}{2\cdot \left \| g_1 \right \|^2 })\cdot g_1 + (1 -
+          \frac{g_1^{T}\cdot g_2}{2\cdot \left \| g_2 \right \|^2 })\cdot g_2]  \\
         \end{array}
 
     In this implementation, :math:`g` represents the gradient of the weights,
@@ -477,8 +479,8 @@ class AdaSumByDeltaWeightWrapCell(Cell):
     .. math::
         \begin{array}{ll}
           w_{t+1}=w_{t} - \alpha \cdot Adasum(g_{1}, g_{2})  \\
-          w_{t+1}=w_{t} - \alpha \cdot [(1 - \frac{g_2^{T}\cdot g_1}{2\cdot \left \| g_1 \right \|^2 })\cdot g_1 +
-           (1 - \frac{g_1^{T}\cdot g_2}{2\cdot \left \| g_2 \right \|^2 })\cdot g_2]  \\
+          w_{t+1}=w_{t} - \alpha \cdot [(1 - \frac{g_2^{T}\cdot g_1}{2\cdot \left \| g_1 \right \|^2 })\cdot g_1 + (1 -
+          \frac{g_1^{T}\cdot g_2}{2\cdot \left \| g_2 \right \|^2 })\cdot g_2]  \\
         \end{array}
 
     In this implementation, :math:`g` represents the weight difference before and after the updating of optimizer,

mindspore/nn/optim/asgd.py

@@ -17,11 +17,11 @@ from __future__ import absolute_import
 
 from mindspore.ops import operations as P
 from mindspore.common.parameter import Parameter
-from mindspore.common.api import ms_function
+from mindspore.common.api import jit
 from mindspore.common.tensor import Tensor
 import mindspore.common.dtype as mstype
 import mindspore
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 
@@ -45,7 +45,7 @@ class ASGD(Optimizer):
 
     :math:`\lambda` represents the decay term, :math:`\mu` and :math:`\eta` are tracked to
     update :math:`ax` and :math:`w`, :math:`t0` represents the point of starting averaging,
-    :math:`\alpha` represents the power for eta update, :math:`ax` represents the averaged
+    :math:`\alpha` represents the power for :math:`\eta` update, :math:`ax` represents the averaged
     parameter value, :math:`t` represents the current step, :math:`g` represents `gradients`,
     :math:`w` represents `params`.
 
@@ -81,7 +81,7 @@ class ASGD(Optimizer):
               If `order_params` in the keys, other keys will be ignored and the element of 'order_params' must be in
               one group of `params`.
 
-        learning_rate (Union[float, int, Tensor, Iterable, LearningRateSchedule]):
+        learning_rate (Union[float, int, Tensor, Iterable, LearningRateSchedule]): learning_rate. Default: 0.1.
 
             - float: The fixed learning rate value. Must be equal to or greater than 0.
 
@@ -96,7 +96,7 @@ class ASGD(Optimizer):
               LearningRateSchedule with step as the input to get the learning rate of current step.
 
         lambd (float): The decay term. Default: 1e-4.
-        alpha (float): The power for eta update. Default: 0.75.
+        alpha (float): The power for :math:`\eta` update. Default: 0.75.
         t0 (float): The point of starting averaging. Default: 1e6.
         weight_decay (Union[float, int, Cell]): Weight decay (L2 penalty). Default: 0.0.
 
@@ -149,7 +149,7 @@ class ASGD(Optimizer):
     """
 
     @opt_init_args_register
-    def __init__(self, params, learning_rate=0.1, lambd=1e-4, alpha=0.75, t0=1e6, weight_decay=0.):
+    def __init__(self, params, learning_rate=0.1, lambd=1e-4, alpha=0.75, t0=1e6, weight_decay=0.0):
 
         super(ASGD, self).__init__(learning_rate, params, weight_decay)
 
@@ -176,7 +176,7 @@ class ASGD(Optimizer):
         self.cast = P.Cast()
         self.squeeze = P.Squeeze()
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         gradients = self.flatten_gradients(gradients)
         gradients = self.decay_weight(gradients)

mindspore/nn/optim/ftrl.py

@@ -16,9 +16,9 @@
 from __future__ import absolute_import
 
 from mindspore.ops import functional as F, composite as C, operations as P
-from mindspore.common.api import ms_function
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
+from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
+from mindspore.common.api import jit
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 from mindspore.nn.optim._dist_optimizer_registry import _register_dist_optimizer
@@ -47,6 +47,62 @@ def _tensor_run_opt_with_sparse_dist(opt, spars_opt, push, pull, l1, l2, lr_powe
     return success
 
 
+def _apply_map_tensor_ftrl(l1, l2, lr_power, learning_rate, linear, weight, moment, indices, values):
+    """Apllpy ftrl optimizer for map parameter"""
+    success = True
+    linear_slice = linear.get(indices)
+    moment_slice = moment.get(indices)
+    weight_slice = weight.get(indices)
+
+    op_pow = P.Pow()
+    op_sign = P.Sign()
+    op_greater = P.Greater()
+    op_select = P.Select()
+    op_abs = P.Abs()
+
+    lr_power_val = -lr_power
+    accu_pow = op_pow(moment_slice, lr_power_val)
+    moment_slice = F.depend(moment_slice, accu_pow)
+    cur_accu = moment_slice + values * values
+    cur_accu_pow = op_pow(cur_accu, lr_power_val)
+    sigma = (cur_accu_pow - accu_pow) / learning_rate
+
+    linear_slice = linear_slice + values - sigma * weight_slice
+
+    update_weight_cond = op_greater(op_abs(linear_slice), l1)
+    updated_weight = (l1 * op_sign(linear_slice) - linear_slice) / (cur_accu_pow / learning_rate + 2 * l2)
+    zeros = zeros_like(weight_slice)
+
+    weight_slice = op_select(update_weight_cond, updated_weight, zeros)
+    moment_slice = cur_accu
+
+    moment.put(indices, moment_slice)
+    linear.put(indices, linear_slice)
+    weight.put(indices, weight_slice)
+
+    return success
+
+
+@_ftrl_opt.register("Function", "Function", "Function", "Function", "Number", "Number", "Number", "Tensor", "MapTensor",
+                    "MapTensor", "MapTensor", "MapTensor", "Bool", "Bool",
+                    "Function", "Bool", "Function", "Bool")
+def _run_map_tensor_opt_with_sparse_dist(opt, spars_opt, push, pull, l1, l2, lr_power, learning_rate, linear,
+                                         gradient, weight, moment, ps_parameter, cache_enable,
+                                         distributed_opt, use_flag, distributed_sparse_opt, use_sparse_flag):
+    """Apply sparse ftrl optimizer to the weight parameter when the gradient is sparse."""
+    success = True
+    indices, values = gradient.get_data()
+    if use_sparse_flag:
+        # PS Mode.
+        success = F.depend(success, distributed_sparse_opt(weight, moment, linear, values, indices))
+    elif cache_enable:
+        # PS Cache mode.
+        _apply_map_tensor_ftrl(l1, l2, lr_power, learning_rate, linear, weight, moment, indices, values)
+    else:
+        raise Exception("Unexpected mode for distributed optimizer.")
+    return success
+
+
 @_ftrl_opt.register("Function", "Function", "Function", "Function", "Number", "Number", "Number", "Tensor", "Tensor",
                     "Tensor", "Tensor", "Tensor", "Bool", "Bool",
                     "Function", "Bool", "Function", "Bool")
@@ -83,6 +139,17 @@ def _tensor_run_opt_with_sparse(opt, spars_opt, push, pull, l1, l2, lr_power, le
     return success
 
 
+@_ftrl_opt.register("Function", "Function", "Function", "Function", "Number", "Number", "Number", "Tensor", "MapTensor",
+                    "MapTensor", "MapTensor", "MapTensor", "Bool", "Bool")
+def _run_map_tensor_opt_with_sparse(opt, spars_opt, push, pull, l1, l2, lr_power, learning_rate, linear,
+                                    gradient, weight, moment, ps_parameter, cache_enable):
+    """Apply sparse ftrl optimizer to the weight parameter when the gradient is sparse."""
+    success = True
+    indices, values = gradient.get_data()
+    _apply_map_tensor_ftrl(l1, l2, lr_power, learning_rate, linear, weight, moment, indices, values)
+    return success
+
+
 @_ftrl_opt.register("Function", "Function", "Function", "Function", "Number", "Number", "Number", "Tensor", "Tensor",
                     "Tensor", "Tensor", "Tensor", "Bool", "Bool")
 def _tensor_run_opt(opt, spars_opt, push, pull, l1, l2, lr_power, learning_rate, linear,
@@ -98,19 +165,22 @@ def _tensor_run_opt(opt, spars_opt, push, pull, l1, l2, lr_power, learning_rate,
     return success
 
 
-def _check_param(initial_accum, lr_power, l1, l2, use_locking, prim_name=None):
+def _check_param(initial_accum, learning_rate, lr_power, l1, l2, use_locking, prim_name=None):
     """Check param."""
     validator.check_value_type("initial_accum", initial_accum, [float], prim_name)
-    validator.check_number("initial_accum", initial_accum, 0.0, Rel.GE, prim_name)
+    validator.check_number("initial_accum", initial_accum, 0.0, validator.GE, prim_name)
+
+    validator.check_value_type("learning_rate", learning_rate, [float], prim_name)
+    validator.check_positive_float(learning_rate, "learning_rate", prim_name)
 
     validator.check_value_type("lr_power", lr_power, [float], prim_name)
-    validator.check_number("lr_power", lr_power, 0.0, Rel.LE, prim_name)
+    validator.check_number("lr_power", lr_power, 0.0, validator.LE, prim_name)
 
     validator.check_value_type("l1", l1, [float], prim_name)
-    validator.check_number("l1", l1, 0.0, Rel.GE, prim_name)
+    validator.check_number("l1", l1, 0.0, validator.GE, prim_name)
 
     validator.check_value_type("l2", l2, [float], prim_name)
-    validator.check_number("l2", l2, 0.0, Rel.GE, prim_name)
+    validator.check_number("l2", l2, 0.0, validator.GE, prim_name)
 
     validator.check_value_type("use_locking", use_locking, [bool], prim_name)
 
@@ -254,7 +324,7 @@ class FTRL(Optimizer):
             raise ValueError(f"For 'FTRL', dynamic learning rate and group learning rate are currently not supported "
                              f"in FTRL, they should all be false, but got dynamic learning rate {self.dynamic_lr} and"
                              f" group learning rate {self.is_group_lr}.")
-        _check_param(initial_accum, lr_power, l1, l2, use_locking, self.cls_name)
+        _check_param(initial_accum, learning_rate, lr_power, l1, l2, use_locking, self.cls_name)
         self.moments = self._parameters.clone(prefix="moments", init=initial_accum)
         self.linear = self._parameters.clone(prefix="linear", init='zeros')
         self.l1 = l1
@@ -276,7 +346,7 @@ class FTRL(Optimizer):
 
         self._init_distributed_opts(use_locking, learning_rate, l1, l2, lr_power)
 
-    @ms_function
+    @jit
     def construct(self, grads):
         params = self._parameters
         moments = self.moments
@@ -316,7 +386,7 @@ class FTRL(Optimizer):
 
         if value == 'CPU':
             self.sparse_opt = P.FusedSparseFtrl(self.lr, self.l1, self.l2, self.lr_power, self.use_locking)
-            self.sparse_opt.add_prim_attr("primitive_target", "CPU")
+            self.sparse_opt.set_device("CPU")
         else:
             self.sparse_opt = P.SparseApplyFtrl(self.lr, self.l1, self.l2, self.lr_power, self.use_locking)
 

mindspore/nn/optim/lamb.py

@@ -22,9 +22,8 @@ from mindspore.ops import composite as C
 from mindspore.ops import functional as F
 from mindspore.ops.operations import _inner_ops as inner
 from mindspore.common.tensor import Tensor
-from mindspore.common.api import ms_function
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
+from mindspore.common.api import jit
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 
@@ -69,8 +68,8 @@ def _check_param_value(beta1, beta2, eps, prim_name):
     validator.check_value_type("beta1", beta1, [float], prim_name)
     validator.check_value_type("beta2", beta2, [float], prim_name)
     validator.check_value_type("eps", eps, [float], prim_name)
-    validator.check_float_range(beta1, 0.0, 1.0, Rel.INC_NEITHER, "beta1", prim_name)
-    validator.check_float_range(beta2, 0.0, 1.0, Rel.INC_NEITHER, "beta2", prim_name)
+    validator.check_float_range(beta1, 0.0, 1.0, validator.INC_NEITHER, "beta1", prim_name)
+    validator.check_float_range(beta2, 0.0, 1.0, validator.INC_NEITHER, "beta2", prim_name)
     validator.check_positive_float(eps, "eps", prim_name)
 
 
@@ -113,7 +112,7 @@ class Lamb(Optimizer):
              /(\sqrt{\hat{\boldsymbol{v}}_{t}}+\epsilon) \\
             &\hspace{5mm}\boldsymbol{w}_{t+1}^{(i)}=\boldsymbol{w}_{t}^{(i)}- \gamma_{t}
              \frac{\boldsymbol{\phi}\left(\left\|\boldsymbol{w}_{t}^{(i)}\right\|\right)}
-             {\left\|\boldsymbol{w}_{t}^{(i)}+\lambda \boldsymbol{w}_{t}^{(i)}\right\|}\left(\boldsymbol{r}_{t}^{(i)}+
+             {\left\|\boldsymbol{r}_{t}^{(i)}+\lambda \boldsymbol{w}_{t}^{(i)}\right\|}\left(\boldsymbol{r}_{t}^{(i)}+
              \lambda \boldsymbol{w}_{t}^{(i)}\right) \\
             &\textbf{end for} \\[-1.ex]
             &\newline
@@ -133,7 +132,7 @@ class Lamb(Optimizer):
         There is usually no connection between a optimizer and mixed precision. But when `FixedLossScaleManager` is used
         and `drop_overflow_update` in `FixedLossScaleManager` is set to False, optimizer needs to set the 'loss_scale'.
         As this optimizer has no argument of `loss_scale`, so `loss_scale` needs to be processed by other means. Refer
-        document `LossScale <https://www.mindspore.cn/tutorials/experts/zh-CN/r1.10/others/mixed_precision.html>`_ to
+        document `LossScale <https://www.mindspore.cn/tutorials/en/r2.0/advanced/mixed_precision.html>`_ to
         process `loss_scale` correctly.
 
         If parameters are not grouped, the `weight_decay` in optimizer will be applied on the network parameters without
@@ -259,7 +258,7 @@ class Lamb(Optimizer):
         self.moments2 = self.params.clone(prefix="lamb_v", init='zeros')
         self.device_ascend = context.get_context("device_target") == "Ascend"
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         weight_decay = self.get_weight_decay()
         lr = self.get_lr()

mindspore/nn/optim/lars.py

@@ -18,9 +18,9 @@ from __future__ import absolute_import
 from mindspore.ops import operations as P
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.common import Tensor, Parameter, dtype as mstype
-from mindspore.common.api import ms_function
+from mindspore.common.api import jit
 from mindspore.nn.optim.optimizer import _grad_scale, Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 
@@ -83,7 +83,7 @@ class LARS(Optimizer):
         \end{array}
 
     :math:`w` represents the network parameters, :math:`g` represents `gradients`,
-    :math:`t` represents the current step, :math:`\delta` represents `weight_decay` in `optimizer`,
+    :math:`t` represents the current step, :math:`\lambda` represents `weight_decay` in `optimizer`,
     :math:`\gamma` represents `learning_rate` in `optimizer`, :math:`\eta` represents `coefficient`.
 
     Args:
@@ -171,7 +171,7 @@ class LARS(Optimizer):
 
         return lr
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         params = self.parameters
         gradients = self.flatten_gradients(gradients)

mindspore/nn/optim/lazyadam.py

@@ -17,17 +17,17 @@ from __future__ import absolute_import
 
 from mindspore.common import dtype as mstype
 from mindspore.common.initializer import initializer
-from mindspore.common.api import ms_function
+from mindspore.common.api import jit
 from mindspore.ops import operations as P
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 from mindspore.nn.optim._dist_optimizer_registry import _register_dist_optimizer
+from mindspore.common._decorator import deprecated
 
 _lazy_adam_opt = C.MultitypeFuncGraph("lazy_adam_opt")
 
@@ -85,6 +85,46 @@ def _run_opt_with_sparse_dist(opt, sparse_opt, push, pull, use_locking, use_nest
     return success
 
 
+@_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "MapTensor", "MapTensor", "MapTensor", "MapTensor",
+                         "Bool", "Bool", "Function", "Bool", "Function", "Bool")
+def _run_map_tensor_opt_with_sparse_dist(opt, sparse_opt, push, pull, use_locking, use_nesterov, target, beta1_power,
+                                         beta2_power, beta1, beta2, eps, lr, gradient, params, m, v,
+                                         ps_parameter, cache_enable, distributed_opt, use_flag, distributed_sparse_opt,
+                                         use_sparse_flag):
+    """Apply sparse lazy adam optimizer to the weight parameter when the gradient is sparse."""
+    success = True
+    indices, values = gradient.get_data()
+    if use_sparse_flag:
+        # PS Mode.
+        success = F.depend(success, distributed_sparse_opt(params, m, v, beta1_power, beta2_power, lr, beta1, beta2,
+                                                           eps, values, indices))
+    else:
+        # PS Cache mode.
+        op_sqrt = P.Sqrt()
+
+        m_slice = m.get(indices)
+        v_slice = v.get(indices)
+
+        next_m = m_slice * beta1 + values * (1 - beta1)
+        next_v = v_slice * beta2 + values * values * (1 - beta2)
+
+        lr_t = lr * op_sqrt(1 - beta2_power) / (1 - beta1_power)
+
+        if use_nesterov:
+            m_temp = beta1 * next_m + values * (1 - beta1)
+            param_update = m_temp / (op_sqrt(next_v) + eps)
+        else:
+            param_update = next_m / (op_sqrt(next_v) + eps)
+
+        params_need_update = params.get(indices)
+        params.put(indices, params_need_update - lr_t * param_update)
+        m.put(indices, next_m)
+        v.put(indices, next_v)
+
+    return success
+
+
 @_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
                          "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool",
                          "Function", "Bool", "Function", "Bool")
@@ -155,6 +195,40 @@ def _run_opt_with_sparse(opt, sparse_opt, push, pull, use_locking, use_nesterov,
     return success
 
 
+@_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "MapTensor", "MapTensor", "MapTensor", "MapTensor",
+                         "Bool", "Bool")
+def _run_map_tensor_opt_with_sparse(opt, sparse_opt, push, pull, use_locking, use_nesterov, target, beta1_power,
+                                    beta2_power, beta1, beta2, eps, lr, gradient, params, m, v, ps_parameter,
+                                    cache_enable):
+    """Apply sparse lazy adam optimizer to the weight parameter when the gradient is sparse(MapTensor)."""
+    success = True
+    indices, values = gradient.get_data()
+
+    op_sqrt = P.Sqrt()
+
+    m_slice = m.get(indices)
+    v_slice = v.get(indices)
+
+    next_m = m_slice * beta1 + values * (1 - beta1)
+    next_v = v_slice * beta2 + values * values * (1 - beta2)
+
+    lr_t = lr * op_sqrt(1 - beta2_power) / (1 - beta1_power)
+
+    if use_nesterov:
+        m_temp = beta1 * next_m + values * (1 - beta1)
+        param_update = m_temp / (op_sqrt(next_v) + eps)
+    else:
+        param_update = next_m / (op_sqrt(next_v) + eps)
+
+    params_need_update = params.get(indices)
+    params.put(indices, params_need_update - lr_t * param_update)
+    m.put(indices, next_m)
+    v.put(indices, next_v)
+
+    return success
+
+
 @_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
                          "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool")
 def _run_opt_with_one_number(opt, sparse_opt, push, pull, use_locking, use_nesterov, target, beta1_power, beta2_power,
@@ -177,8 +251,8 @@ def _check_param_value(beta1, beta2, eps, weight_decay, prim_name):
     validator.check_value_type("beta2", beta2, [float], prim_name)
     validator.check_value_type("eps", eps, [float], prim_name)
     validator.check_value_type("weight_dacay", weight_decay, [float], prim_name)
-    validator.check_float_range(beta1, 0.0, 1.0, Rel.INC_NEITHER, "beta1", prim_name)
-    validator.check_float_range(beta2, 0.0, 1.0, Rel.INC_NEITHER, "beta2", prim_name)
+    validator.check_float_range(beta1, 0.0, 1.0, validator.INC_NEITHER, "beta1", prim_name)
+    validator.check_float_range(beta2, 0.0, 1.0, validator.INC_NEITHER, "beta2", prim_name)
     validator.check_positive_float(eps, "eps", prim_name)
     validator.check_non_negative_float(weight_decay, "weight_decay", prim_name)
 
@@ -332,6 +406,7 @@ class LazyAdam(Optimizer):
         >>> model = ms.Model(net, loss_fn=loss, optimizer=optim)
     """
 
+    @deprecated("2.0", "Adam", False)
     @opt_init_args_register
     def __init__(self, params, learning_rate=1e-3, beta1=0.9, beta2=0.999, eps=1e-8, use_locking=False,
                  use_nesterov=False, weight_decay=0.0, loss_scale=1.0):
@@ -352,14 +427,14 @@ class LazyAdam(Optimizer):
         self.moment2 = self._parameters.clone(prefix="moment2", init='zeros')
         self.opt = P.Adam(use_locking, use_nesterov)
         self.sparse_opt = P.FusedSparseLazyAdam(use_locking, use_nesterov)
-        self.sparse_opt.add_prim_attr("primitive_target", "CPU")
+        self.sparse_opt.set_device("CPU")
         self._ps_pull = P.Pull()
         self._ps_push = P.Push("Adam", [0, 1, 2])
         self._ps_push.add_prim_attr("use_nesterov", use_nesterov)
 
         self._init_distributed_opts(use_locking, use_nesterov)
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         gradients = self.flatten_gradients(gradients)
         gradients = self.decay_weight(gradients)

mindspore/nn/optim/momentum.py

@@ -18,9 +18,9 @@ from __future__ import absolute_import
 from mindspore.ops import functional as F, composite as C, operations as P
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
-from mindspore.common.api import ms_function
+from mindspore.common.api import jit
 import mindspore.common.dtype as mstype
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 from mindspore.nn.optim._dist_optimizer_registry import _register_dist_optimizer
@@ -69,19 +69,20 @@ class Momentum(Optimizer):
     learning <https://dl.acm.org/doi/10.5555/3042817.3043064>`_ for more details.
 
     .. math::
-            v_{t+1} = v_{t} \ast u + grad
+        v_{t+1} = v_{t} \ast u + grad
 
     If use_nesterov is True:
 
     .. math::
-            p_{t+1} =  p_{t} - (grad \ast lr + v_{t+1} \ast u \ast lr)
+        p_{t+1} =  p_{t} - (grad \ast lr + v_{t+1} \ast u \ast lr)
 
     If use_nesterov is False:
 
     .. math::
-            p_{t+1} = p_{t} - lr \ast v_{t+1}
+        p_{t+1} = p_{t} - lr \ast v_{t+1}
 
-    Here: where grad, lr, p, v and u denote the gradients, learning_rate, params, moments, and momentum respectively.
+    Here: where :math:`grad`, :math:`lr`, :math:`p`, :math:`v` and :math:`u` denote the gradients,
+    learning_rate, params, moments, and momentum respectively.
 
     Note:
         If parameters are not grouped, the `weight_decay` in optimizer will be applied on the network parameters without
@@ -208,7 +209,7 @@ class Momentum(Optimizer):
         self._get_distributed_optimizer_list("momentum", use_nesterov=self.use_nesterov)
         self.use_dist_optimizer = self._use_distibuted_optimizer()
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         params = self.params
         moments = self.moments