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

mindspore/dataset/vision/validators.py

@@ -351,6 +351,7 @@ def check_resize_interpolation(method):
 
     return new_method
 
+
 def check_device_target(method):
     """A wrapper that wraps a parameter checker"""
 

mindspore/device_context/ascend/device.py

@@ -66,7 +66,7 @@ def is_available():
 
 def _is_supported():
     device_target = ms.context.get_context("device_target")
-    if device_target == 'CPU' or device_target == 'GPU':
+    if device_target in ['CPU', 'GPU']:
         logger.error(f"{device_target} device is not supported. Please use correct device")
         return False
     return True

mindspore/device_context/ascend/op_tuning.py

@@ -21,7 +21,7 @@ except ImportError:
 from .device import _is_supported
 
 function_status = {'op_compile': False, 'aoe_tune_mode': False,
-                   'aoe_job_type': False}
+                   'aoe_job_type': False, 'aclnn_cache': False}
 
 
 def op_compile(value):
@@ -121,3 +121,37 @@ def aoe_job_type(config):
             f"For 'aoe_job_type', the config must be one of {aoe_cfgs}, but got {config}."
         )
     AscendOpTuningConf.get_instance().set_aoe_job_type(config)
+
+
+def aclnn_cache(enable_global_cache=False, cache_queue_length=10000):
+    """
+    Configure aclnn cache parameters.
+
+    Args:
+        enable_global_cache (bool): Set the calnn cache to global when GRAPH_MODE.
+            Default: ``False``.
+        cache_queue_length (int, optional): Set the cache queue length.
+            Default: ``10000``.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.device_context.ascend.op_tuning.aclnn_cache(True, 10000)
+    """
+    if not function_status['aclnn_cache']:
+        function_status['aclnn_cache'] = True
+        if not _is_supported():
+            return
+    # Check the configuration environment whether valid
+    if AscendOpTuningConf.get_instance().is_aclnn_cache_configured():
+        raise RuntimeError("The 'aclnn_cache' can not be set repeatedly.")
+    cache_cfgs = [True, False]
+    if enable_global_cache not in cache_cfgs:
+        raise ValueError(
+            f"For 'aclnn_cache', the config must be one of {cache_cfgs}, but got {enable_global_cache}."
+        )
+    AscendOpTuningConf.get_instance().set_aclnn_global_cache(enable_global_cache)
+    if cache_queue_length < 0:
+        raise ValueError(
+            f"For 'aclnn_cache', the config must greater than 0, but got {enable_global_cache}."
+        )
+    AscendOpTuningConf.get_instance().set_cache_queue_length(cache_queue_length)

mindspore/device_context/gpu/__init__.py

@@ -14,8 +14,8 @@
 # ============================================================================
 """init file for GPU device context"""
 
+__all__ = ['device_count', 'is_available']
+
 from .device import device_count, is_available
 from .op_precision import *
 from .op_tuning import *
-
-__all__ = ["device_count", "is_available"]

mindspore/device_context/gpu/device.py

@@ -64,7 +64,7 @@ def is_available():
 
 def _is_supported():
     device_target = ms.context.get_context("device_target")
-    if device_target == 'CPU' or device_target == 'Ascend':
+    if device_target in ['CPU', 'Ascend']:
         logger.error(f"{device_target} device is not supported. Please use correct device")
         return False
     return True

mindspore/device_context/gpu/op_precision.py

@@ -29,7 +29,8 @@ function_status = {'matmul_allow_tf32': False, 'conv_allow_tf32': False}
 def matmul_allow_tf32(value):
     """
     Whether to convert FP32 to TF32 for Matmul operators.
-    For detailed information, please refer to `CUBLAS_COMPUTE_32F_FAST_TF32 <https://docs.nvidia.com/cuda/cublas/index.html>`_.
+    For detailed information, please refer to `CUBLAS_COMPUTE_32F_FAST_TF32
+    <https://docs.nvidia.com/cuda/cublas/index.html>`_.
 
     Args:
         value (bool): Whether to convert FP32 to TF32 for Matmul operators. If not configured, the framework
@@ -50,7 +51,8 @@ def matmul_allow_tf32(value):
 def conv_allow_tf32(value):
     """
     Whether to convert FP32 to TF32 for Conv operators.
-    For detailed information, please refer to `CUBLAS_COMPUTE_32F_FAST_TF32 <https://docs.nvidia.com/cuda/cublas/index.html>`_.
+    For detailed information, please refer to `CUBLAS_COMPUTE_32F_FAST_TF32
+    <https://docs.nvidia.com/cuda/cublas/index.html>`_.
 
     Args:
         value (bool): Whether to convert FP32 to HF32 for Conv operators. If not configured, the framework defaults

mindspore/device_context/gpu/op_tuning.py

@@ -27,7 +27,8 @@ function_status = {'conv_fprop_algo': False, 'conv_wgrad_algo': False, 'conv_dgr
 def conv_fprop_algo(mode):
     """
     Specifies convolution forward algorithm.
-    For detailed information, please refer to `NVIDA cuDNN about cudnnConvolutionForward <https://docs.nvidia.com/deeplearning/cudnn/latest/api/cudnn-cnn-library.html>`_.
+    For detailed information, please refer to `NVIDA cuDNN about cudnnConvolutionForward
+    <https://docs.nvidia.com/deeplearning/cudnn/latest/api/cudnn-cnn-library.html>`_.
 
     Args:
         mode (str): convolution forward algorithm. If not configured, the framework defaults to 'normal'.
@@ -80,7 +81,8 @@ def conv_fprop_algo(mode):
 def conv_wgrad_algo(mode):
     """
     Specifies convolution filter grad algorithm.
-    For detailed information, please refer to `NVIDA cuDNN <https://docs.nvidia.com/deeplearning/cudnn/latest/api/cudnn-cnn-library.html>`_.
+    For detailed information, please refer to `NVIDA cuDNN
+    <https://docs.nvidia.com/deeplearning/cudnn/latest/api/cudnn-cnn-library.html>`_.
 
     Args:
         mode (str): convolution filter grad algorithm. If not configured, the framework defaults to 'normal'.
@@ -129,7 +131,8 @@ def conv_wgrad_algo(mode):
 def conv_dgrad_algo(mode):
     """
     Specifies convolution data grad algorithm.
-    For detailed information, please refer to `NVIDA cuDNN <https://docs.nvidia.com/deeplearning/cudnn/latest/api/cudnn-cnn-library.html>`_.
+    For detailed information, please refer to `NVIDA cuDNN
+    <https://docs.nvidia.com/deeplearning/cudnn/latest/api/cudnn-cnn-library.html>`_.
 
     Args:
         mode (str): convolution data grad algorithm. If not configured, the framework defaults to 'normal'.

mindspore/device_manager.py

@@ -15,16 +15,14 @@
 
 """Device manager interfaces."""
 
+__all__ = ['set_device', 'set_deterministic', 'get_current_device']
+
 import os
 from mindspore import log as logger
 from mindspore._c_expression import DeviceManagerConf, DeviceContextManager, MSContext, CollectiveManager
 from mindspore._checkparam import args_type_check
 from mindspore.parallel._ps_context import _need_reset_device_target_for_ps
 
-
-__all__ = ['set_device', 'set_deterministic', 'get_current_device']
-
-
 class DeviceInfo(tuple):
     """
     DeviceInfo class. Store the current device target and the corresponding device id.
@@ -120,14 +118,21 @@ def set_deterministic(deterministic):
     """
     Enables or disables deterministic computing.
 
-    When deterministic computing is enabled, the same output is generated if an operator is executed
-    for multiple times with the same hardware and input.This often slows down operator execution.
-    In distributed scenario, we suggest user to set deterministic mode before
-    calling :func:`mindspore.communication.init` to enable deterministic operation for
-    communication operators in the global communication group.
+    This configuration is a global configuration, and once enabled, subsequent calculation operations
+    will follow the configuration setting. When deterministic computing is enabled, the same output
+    is generated if an operator is executed for multiple times with the same hardware and input. This often
+    slows down operator execution.
 
     The framework not enabled deterministic computation by default.
 
+    Note:
+        - In distributed scenario, we suggest user to set deterministic computing before
+          calling :func:`mindspore.communication.init` to enable deterministic operation for
+          communication operators in the global communication group.
+        - The fixed method for deterministic calculation must be in the same main process as the network,
+          operator, etc. Only one deterministic state can be set in the same thread, and it is not recommended
+          to set deterministic state multiple times in one thread.
+
     Args:
         deterministic (bool): Whether to enable deterministic computing.
 
@@ -139,6 +144,8 @@ def set_deterministic(deterministic):
     if DeviceManagerConf.get_instance().is_deterministic_configured():
         raise RuntimeError("The 'mindspore.set_deterministic' can not be set repeatedly.")
 
+    logger.info(f"Set deterministic setting to '{deterministic}'.")
+
     # Must wait for all async created groups to be initialized so that
     # deterministic feature could be consistent between all processes.
     CollectiveManager.get_instance().wait_all_comm_init()

mindspore/experimental/llm_boost/ascend_native/llama_boost_ascend_native.py

@@ -20,7 +20,6 @@ from mindspore.common import Tensor, dtype
 from mindspore.experimental.llm_boost.ascend_native.llm_boost import LLMBoost
 from mindspore.experimental.llm_boost.register import LlmBoostRegister, LlmBoostType
 
-
 def RoundUp(val: int, align: int) -> int:
     if align == 0:
         return 0
@@ -68,7 +67,7 @@ class LlamaBoostAscendNative(LLMBoost):
     def _prepare_single_layer(self, ckpt, config, id):
         """ prepares the dictionary of weights of a single layer """
         prefix = 'model.layers.' + str(id)
-        is_last = (id == config.num_layers-1)
+        is_last = id == config.num_layers-1
         layer = 'layers.' + str(id) + '.'
         l_dict = {key: value for key, value in ckpt.items() if layer in key}
         if config.n_kv_heads is None:
@@ -96,8 +95,8 @@ class LlamaBoostAscendNative(LLMBoost):
         else:
             raise RuntimeError("hidden size and ffn hidden size must be divided by rank size without remainder.  \
                                 hidden_size: ", hid_size, " ffn_hidden_size: ", ffn_hid, " rank_size: ", rank_size)
-        quant = (self._get_from_dict(l_dict, "_weight_quantizer") is not None)
-        unite_qkv = (config.num_heads == config.n_kv_heads)
+        quant = self._get_from_dict(l_dict, "_weight_quantizer") is not None
+        unite_qkv = config.num_heads == config.n_kv_heads
         self.dictionary[prefix + ".attention_norm.weight"] = \
             Tensor(self._get_from_dict(l_dict, "attention_norm"), dtype=dtype.float16)
         self.dictionary[prefix + ".ffn_norm.weight"] = \

mindspore/experimental/llm_boost/atb/boost_base.py

@@ -18,7 +18,6 @@ import numpy as np
 import mindspore as ms
 from mindspore import ops, Tensor
 from mindspore import log as logger
-from mindspore.ops import operations as P
 import mindspore.common.dtype as mstype
 from mindspore._c_expression import _set_format
 from mindspore.common.parameter import Parameter
@@ -95,8 +94,8 @@ class AtbBoostBase:
             self.max_base_len, dtype=self.dtype, need_nz=self.need_nz
         )
 
-        self.cast = P.Cast()
-        self.reshape = P.Reshape()
+        self.cast = ops.Cast()
+        self.reshape = ops.Reshape()
         self.kv_quant = None
         self.rank_id = get_real_rank()
         self.device_num = get_real_group_size()

mindspore/experimental/optim/adadelta.py

@@ -1,34 +1,24 @@
-# Copyright 2023 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/adadelta.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adadelta"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 import mindspore.common.dtype as mstype
 from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_than, check_not_less_than_without_equal
 from mindspore import _checkparam as validator
 from mindspore import jit
 
-_adadelta_opt = C.MultitypeFuncGraph("adadelta_opt")
+_adadelta_opt = ops.MultitypeFuncGraph("adadelta_opt")
 
 
 @_adadelta_opt.register("Function", "Number", "Number", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
 def _tensor_run_opt(opt, rho, epsilon, learning_rate, weight, accum, accum_update, gradient):
     """Apply adadelta optimizer to the weight parameter."""
     success = True
-    success = F.depend(success, opt(weight, accum, accum_update, learning_rate, rho, epsilon, gradient))
+    success = ops.depend(success, opt(weight, accum, accum_update, learning_rate, rho, epsilon, gradient))
     return success
 
 
@@ -60,6 +50,9 @@ class Adadelta(Optimizer):
             &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
         \end{aligned}
 
+    For more details about Adadelta algorithm, please refer to `ADADELTA: An Adaptive Learning Rate Method
+    <https://arxiv.org/abs/1212.5701>`_.
+
     .. warning::
         This is an experimental optimizer API that is subject to change.
         This module must be used with lr scheduler module in `LRScheduler Class
@@ -129,18 +122,18 @@ class Adadelta(Optimizer):
 
         self.accum = self.parameters.clone(prefix="accum", init=0)
         self.accum_update = self.parameters.clone(prefix="accum_update", init=0)
-        self.opt = P.ApplyAdadelta()
-        self.op_cast = P.Cast()
+        self.opt = ops.ApplyAdadelta()
+        self.op_cast = ops.Cast()
 
     @jit
     def implementation(self, lr, rho, eps, maximize, weight_decay, start_id, end_id, gradients):
         """Extract the common computing part for acceleration"""
         params = self.parameters[start_id: end_id]
-        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = tuple([grad if not maximize else ops.neg(grad) for grad in gradients[start_id: end_id]])
         grads = self._decay_weight(weight_decay, params, grads)
         accum = self.accum[start_id: end_id]
         accum_update = self.accum_update[start_id: end_id]
-        self.hyper_map(F.partial(_adadelta_opt, self.opt, rho, eps, lr),
+        self.hyper_map(ops.partial(_adadelta_opt, self.opt, rho, eps, lr),
                        params, accum, accum_update, grads)
         return True
 

mindspore/experimental/optim/adagrad.py

@@ -1,34 +1,24 @@
-# Copyright 2023 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/adagrad.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adagrad"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 from mindspore.common import Tensor, Parameter
 import mindspore.common.dtype as mstype
 from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_than, check_not_less_than_without_equal
 from mindspore import jit
 
-_adagrad_opt = C.MultitypeFuncGraph("adagrad_opt")
+_adagrad_opt = ops.MultitypeFuncGraph("adagrad_opt")
 
 
 @_adagrad_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor")
 def _tensor_run_opt(opt, learning_rate, weight, accum, gradient):
     """Apply adagrad optimizer to the weight parameter."""
     success = True
-    success = F.depend(success, opt(weight, accum, learning_rate, gradient))
+    success = ops.depend(success, opt(weight, accum, learning_rate, gradient))
     return success
 
 
@@ -57,6 +47,9 @@ class Adagrad(Optimizer):
             &\rule{160mm}{0.4pt}                                                          \\[-1.ex]
        \end{aligned}
 
+    For more details about Adagrad algorithm, please refer to `Adaptive Subgradient Methods for Online Learning and \
+    Stochastic Optimization <https://jmlr.org/papers/v12/duchi11a.html>`_.
+
     .. warning::
         This is an experimental optimizer API that is subject to change.
         This module must be used with lr scheduler module in `LRScheduler Class
@@ -129,22 +122,22 @@ class Adagrad(Optimizer):
         super(Adagrad, self).__init__(params, defaults)
 
         self.accum = self.parameters.clone(prefix="accum", init=initial_accumulator_value)
-        self.op_cast = P.Cast()
+        self.op_cast = ops.Cast()
         self.step_t = Parameter(Tensor(0, mstype.int32), "step_t")
         self.increase_tensor = Tensor(1, mstype.int32)
-        self.assignadd = P.AssignAdd()
-        self.assign = P.Assign()
+        self.assignadd = ops.AssignAdd()
+        self.assign = ops.Assign()
 
     @jit
     def implementation(self, eps, lr, lr_decay, maximize, weight_decay, start_id, end_id, gradients):
         """Extract the common computing part for acceleration"""
-        opt = P.ApplyAdagradV2(epsilon=eps, update_slots=True)
+        opt = ops.ApplyAdagradV2(epsilon=eps, update_slots=True)
         decay_lr = lr / (1 + self.step_t * lr_decay)
         params = self.parameters[start_id: end_id]
-        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = tuple([grad if not maximize else ops.neg(grad) for grad in gradients[start_id: end_id]])
         grads = self._decay_weight(weight_decay, params, grads)
         accum = self.accum[start_id: end_id]
-        self.hyper_map(F.partial(_adagrad_opt, opt, decay_lr), params, accum, grads)
+        self.hyper_map(ops.partial(_adagrad_opt, opt, decay_lr), params, accum, grads)
         return True
 
     def construct(self, gradients):

mindspore/experimental/optim/adam.py

@@ -1,29 +1,19 @@
-# Copyright 2023 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/adam.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adam"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 import mindspore.common.dtype as mstype
 from mindspore.experimental.optim.optimizer import Optimizer
 from mindspore.common.api import jit
 
-_adam_opt = C.MultitypeFuncGraph("adam_opt")
-adam_op = P.Adam(False, False)
+_adam_opt = ops.MultitypeFuncGraph("adam_opt")
+adam_op = ops.Adam(False, False)
 
 
 @_adam_opt.register("Tensor", "Tensor", "Float", "Float", "Float", "Tensor",
@@ -81,6 +71,9 @@ class Adam(Optimizer):
             &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
        \end{aligned}
 
+    For more details about Adam algorithm, please refer to `Adam: A Method for Stochastic Optimization
+    <https://arxiv.org/abs/1412.6980>`_.
+
     .. warning::
         The implementation formula of this optimizer interface is not completely consistent with that in the paper.
         If you want to use an interface that is completely consistent, it is recommended to use
@@ -160,12 +153,12 @@ class Adam(Optimizer):
         self.max_exp_avg_sq = self.parameters.clone(prefix="max_exp_avg_sq", init='zeros')
         self.state_step = Parameter(Tensor(0, mstype.int32), "state_step")
         self.increase_tensor = Tensor(1, mstype.int32)
-        self.assignadd = P.AssignAdd()
-        self.op_add = P.AddN()
-        self.op_mul = P.Mul()
-        self.op_pow = P.Pow()
-        self.adam_opt = P.Adam(False, False)
-        self.op_cast = P.Cast()
+        self.assignadd = ops.AssignAdd()
+        self.op_add = ops.AddN()
+        self.op_mul = ops.Mul()
+        self.op_pow = ops.Pow()
+        self.adam_opt = ops.Adam(False, False)
+        self.op_cast = ops.Cast()
 
     @jit
     def implementation(self, beta1, beta2, eps, lr, start_id, end_id, gradients, maximize, weight_decay):
@@ -173,9 +166,9 @@ class Adam(Optimizer):
         beta1_power = self.op_pow(beta1, self.state_step)
         beta2_power = self.op_pow(beta2, self.state_step)
         params = self.parameters[start_id: end_id]
-        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = tuple([grad if not maximize else ops.neg(grad) for grad in gradients[start_id: end_id]])
         grads = self._decay_weight(weight_decay, params, grads)
-        self.hyper_map(F.partial(_adam_opt, beta1_power, beta2_power, beta1, beta2, eps, lr),
+        self.hyper_map(ops.partial(_adam_opt, beta1_power, beta2_power, beta1, beta2, eps, lr),
                        grads, params,
                        self.exp_avg[start_id: end_id], self.exp_avg_sq[start_id: end_id])
         return True

mindspore/experimental/optim/adamax.py

@@ -1,21 +1,10 @@
-# Copyright 2023 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/adamax.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adamax"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
 from mindspore.common import Tensor, Parameter
 import mindspore.common.dtype as mstype
 from mindspore import _checkparam as validator
@@ -23,17 +12,17 @@ from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_tha
 from mindspore import ops
 from mindspore import jit
 
-_adamax_opt = C.MultitypeFuncGraph("adamax_opt")
+_adamax_opt = ops.MultitypeFuncGraph("adamax_opt")
 
 
 @_adamax_opt.register("Number", "Number", "Number", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
 def _tensor_run_opt(beta1, beta2, eps, clr, param, grad, exp_avg, exp_inf):
     """Apply adamax optimizer to the weight parameter."""
-    F.assign(exp_avg, exp_avg * beta1 + grad * (1-beta1))
+    ops.assign(exp_avg, exp_avg * beta1 + grad * (1-beta1))
     norm_buf = ops.cat([ops.unsqueeze(exp_inf * beta2, 0), ops.unsqueeze(grad.abs().add(eps), 0)], 0)
-    F.assign(exp_inf, ops.amax(norm_buf, 0))
+    ops.assign(exp_inf, ops.amax(norm_buf, 0))
 
-    F.assign(param, param - clr * exp_avg / exp_inf)
+    ops.assign(param, param - clr * exp_avg / exp_inf)
     return True
 
 
@@ -68,6 +57,9 @@ class Adamax(Optimizer):
         This module must be used with lr scheduler module in `LRScheduler Class
         <https://www.mindspore.cn/docs/en/master/api_python/mindspore.experimental.html#lrscheduler-class>`_ .
 
+    For more details about Adamax algorithm, please refer to `Adam: A Method for Stochastic Optimization
+    <https://arxiv.org/abs/1412.6980>`_.
+
     Args:
         params (Union[list(Parameter), list(dict)]): list of parameters to optimize or dicts defining
             parameter groups.
@@ -135,8 +127,8 @@ class Adamax(Optimizer):
         self.exp_avg = self.parameters.clone(prefix="exp_avg", init='zeros')
         self.exp_inf = self.parameters.clone(prefix="exp_inf", init='zeros')
         self.increase_tensor = Tensor(1, mstype.int32)
-        self.assignadd = P.AssignAdd()
-        self.op_cast = P.Cast()
+        self.assignadd = ops.AssignAdd()
+        self.op_cast = ops.Cast()
 
     @jit
     def implementation(self, group_id, lr, gradients, maximize, weight_decay, beta1, beta2, eps):
@@ -144,13 +136,13 @@ class Adamax(Optimizer):
         start_id = self.group_start_id[group_id]
         end_id = self.group_start_id[group_id + 1]
         params = self.parameters[start_id: end_id]
-        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = tuple([grad if not maximize else ops.neg(grad) for grad in gradients[start_id: end_id]])
         grads = self._decay_weight(weight_decay, params, grads)
         exp_avg = self.exp_avg[start_id: end_id]
         exp_inf = self.exp_inf[start_id: end_id]
         bias_correction = 1 - beta1 ** self.step_t
         clr = lr / bias_correction
-        self.hyper_map(F.partial(_adamax_opt, beta1, beta2, eps, clr),
+        self.hyper_map(ops.partial(_adamax_opt, beta1, beta2, eps, clr),
                        params, grads, exp_avg, exp_inf)
         return True