class_docstr: LAMB Optimizer.
  
  Referenced from 'Large Batch Optimization for Deep Learning: Training BERT in 76 minutes'
  (https://arxiv.org/pdf/1904.00962.pdf)
  
  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
  bias_correction : bool, default True
      Whether or not to apply bias correction
  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.
class_docstr_tok: ["LAMB","Optimizer",".","Referenced","from","'Large","Batch","Optimization","for","Deep","Learning",":","Training","BERT","in","76","minutes'","(","https",":","\/\/arxiv.org\/pdf\/1904.00962.pdf",")","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","bias_correction",":","bool",",","default","True","Whether","or","not","to","apply","bias","correction","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()`.
docstr_tok: ["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 fused_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()`.
      """
      aggregate = self.aggregate_num > 1
      for weight, grad in zip(weights, grads):
          aggregate = (
              aggregate
              and weight.stype == "default"
              and grad.stype == "default"
          )
      self._update_count(indices)
      lrs = self._get_lrs(indices)
      wds = self._get_wds(indices)
  
      if aggregate:
          kwargs = {
              "beta1": self.beta1,
              "beta2": self.beta2,
              "epsilon": self.epsilon,
              "bias_correction": self.bias_correction,
              "rescale_grad": self.rescale_grad,
          }
          if self.clip_gradient:
              kwargs["clip_gradient"] = self.clip_gradient
          if self.lower_bound:
              kwargs["lower_bound"] = self.lower_bound
          if self.upper_bound:
              kwargs["upper_bound"] = self.upper_bound
  
          step_counts = []
          for index in indices:
              step_counts.append(self._index_update_count[index])
  
          multi_precision = (
              self.multi_precision and weights[0].dtype == numpy.float16
          )
  
          if not multi_precision:
              mean, var = list(zip(*states))
              multi_lamb_update(
                  weights,
                  grads,
                  mean,
                  var,
                  out=weights,
                  step_count=step_counts,
                  lrs=lrs,
                  wds=wds,
                  **kwargs,
              )
          else:
              weights32, mean_var = list(zip(*states))
              mean, var = list(zip(*mean_var))
              multi_mp_lamb_update(
                  weights,
                  grads,
                  mean,
                  var,
                  weights32,
                  out=weights,
                  step_count=step_counts,
                  lrs=lrs,
                  wds=wds,
                  **kwargs,
              )
      else:
          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]
              kwargs = {
                  "beta1": self.beta1,
                  "beta2": self.beta2,
                  "epsilon": self.epsilon,
                  "bias_correction": self.bias_correction,
                  "rescale_grad": self.rescale_grad,
                  "t": t,
              }
              if self.clip_gradient:
                  kwargs["clip_gradient"] = self.clip_gradient
  
              multi_precision = (
                  self.multi_precision and weight.dtype == numpy.float16
              )
  
              if multi_precision:
                  weight32 = state[0]
                  mean, var = state[1]
                  g = mp_lamb_update_phase1(
                      weight, grad, mean, var, weight32, wd=wd, **kwargs
                  )
  
                  kwargs = {}
                  if self.lower_bound:
                      kwargs["lower_bound"] = self.lower_bound
                  if self.upper_bound:
                      kwargs["upper_bound"] = self.upper_bound
                  r_1 = weight32.norm()
                  r_2 = g.norm()
                  mp_lamb_update_phase2(
                      weight,
                      g,
                      r_1,
                      r_2,
                      weight32,
                      lr=lr,
                      out=weight,
                      **kwargs,
                  )
              else:
                  mean, var = state
                  g = lamb_update_phase1(
                      weight, grad, mean, var, wd=wd, **kwargs
                  )
  
                  kwargs = {}
                  if self.lower_bound:
                      kwargs["lower_bound"] = self.lower_bound
                  if self.upper_bound:
                      kwargs["upper_bound"] = self.upper_bound
                  r_1 = weight.norm()
                  r_2 = g.norm()
                  lamb_update_phase2(
                      weight, g, r_1, r_2, lr=lr, out=weight, **kwargs
                  )
  
file_path: mxnet/python/mxnet/optimizer/lamb.py
imports: from __future__ import absolute_import
  import numpy
  from ..ndarray import zeros, clip, sqrt, where, square, ones_like, maximum, minimum
  from ..ndarray import lamb_update_phase1, lamb_update_phase2, mp_lamb_update_phase1, mp_lamb_update_phase2
  from ..ndarray.contrib import multi_lamb_update, multi_mp_lamb_update
  from .optimizer import Optimizer, register
label_desc: Use image node_id 4 for calling the LAMB obj's underlying member method code with example usage: obj.fused_step(indices, weights, grads, states) without return types