class_docstr: LANS Optimizer.
  
  Referenced from 'Accelerated Large Batch Optimization of BERT Pretraining in 54 minutes'
  (http://arxiv.org/abs/2006.13484)
  
  Parameters
  ----------
  learning_rate : float, default 0.001
      The initial learning rate. If None, the optimization will use the
      learning rate from ``lr_scheduler``. If not None, it will overwrite
      the learning rate in ``lr_scheduler``. If None and ``lr_scheduler``
      is also None, then it will be set to 0.01 by default.
  beta1 : float, default 0.9
      Exponential decay rate for the first moment estimates.
  beta2 : float, default 0.999
      Exponential decay rate for the second moment estimates.
  epsilon : float, default 1e-6
      Small value to avoid division by 0.
  lower_bound : float, default None
      Lower limit of norm of weight
  upper_bound : float, default None
      Upper limit of norm of weight
  aggregate_num : int, default 4
      Number of weights to be aggregated in a list.
      They are passed to the optimizer for a single optimization step.
      In default, all the weights are aggregated.
  use_fused_step : bool, default True
      Whether or not to use fused kernels for optimizer.
      When use_fused_step=False, step is called,
      otherwise, fused_step is called.
docstr: Perform a fused optimization step using gradients and states.
  Fused kernel is used for update.
  
  Parameters
  ----------
  indices : list of int
      List of unique indices of the parameters into the individual learning rates
      and weight decays. Learning rates and weight decay may be set via `set_lr_mult()`
      and `set_wd_mult()`, respectively.
  weights : list of NDArray
      List of parameters to be updated.
  grads : list of NDArray
      List of gradients of the objective with respect to this parameter.
  states : List of any obj
      List of state returned by `create_state()`.
code: def step(self, indices, weights, grads, states):
      """Perform a fused optimization step using gradients and states.
      Fused kernel is used for update.
  
      Parameters
      ----------
      indices : list of int
          List of unique indices of the parameters into the individual learning rates
          and weight decays. Learning rates and weight decay may be set via `set_lr_mult()`
          and `set_wd_mult()`, respectively.
      weights : list of NDArray
          List of parameters to be updated.
      grads : list of NDArray
          List of gradients of the objective with respect to this parameter.
      states : List of any obj
          List of state returned by `create_state()`.
      """
      for index, weight, grad, state in zip(
          indices, weights, grads, states
      ):
          self._update_count(index)
          lr = self._get_lr(index)
          wd = self._get_wd(index)
          t = self._index_update_count[index]
  
          # preprocess grad
          grad *= self.rescale_grad
          grad /= grad.norm()
          if self.clip_gradient is not None:
              grad = clip(grad, -self.clip_gradient, self.clip_gradient)
  
          # update mean, var
          mean, var = state
          mean[:] *= self.beta1
          mean[:] += (1.0 - self.beta1) * grad
          var[:] *= self.beta2
          var[:] += (1.0 - self.beta2) * square(grad)
  
          r1 = weight.norm()
          if self.lower_bound is not None:
              r1 = maximum(r1, self.lower_bound)
          if self.upper_bound is not None:
              r1 = minimum(r1, self.upper_bound)
  
          # apply bias correction
          coef1 = 1.0 - self.beta1**t
          coef2 = 1.0 - self.beta2**t
          mean_hat = mean / coef1
          var_hat = var / coef2
          sqrt(var_hat, out=var_hat)
          var_hat += self.epsilon
          mean_hat /= var_hat
          mean_hat += wd * weight
  
          g = mean_hat
          r2 = g.norm()
  
          # calculate lans_trust_ratio for first part
          ratio_m = r1 / r2
          # becomes NaN if ratio == NaN or 0, otherwise 0
          nan_or_zero = 1 - ratio_m / ratio_m
          r_m = where(nan_or_zero, ones_like(ratio_m), ratio_m)
  
          # update weight using first part of the estimator
          g *= lr * r_m * self.beta1
          weight[:] -= g
  
          # calculate the second part of the estimator
          mean_hat = grad / var_hat
          mean_hat += wd * weight
  
          g = mean_hat
          r2 = g.norm()
  
          # calculate lans_trust_ratio for second part
          ratio_g = r1 / r2
          # becomes NaN if ratio == NaN or 0, otherwise 0
          nan_or_zero = 1 - ratio_g / ratio_g
          r_g = where(nan_or_zero, ones_like(ratio_g), ratio_g)
  
          # update weight using second part of the estimator
          g *= lr * r_g * (1 - self.beta1)
          weight[:] -= g
  
file_path: mxnet/python/mxnet/optimizer/lans.py
imports: from __future__ import absolute_import
  import numpy
  from ..ndarray import zeros, clip, sqrt, where, square, ones_like, maximum, minimum
  from ..ndarray.contrib import multi_lans_update, multi_mp_lans_update
  from .optimizer import Optimizer, register
label_desc: Use image node_id 3 for calling the LANS obj's underlying member method code with example usage: obj.step(indices, weights, grads, states) without return types