class_docstr: Implements Adam algorithm.
  
  This implementation is modified from torch.optim.Adam based on:
  `Fixed Weight Decay Regularization in Adam`
  (see https://arxiv.org/abs/1711.05101)
  
  It has been proposed in `Adam: A Method for Stochastic Optimization`_.
  
  Args:
      params (iterable): iterable of parameters to optimize or dicts defining
          parameter groups
      lr (float, optional): learning rate (default: 1e-3)
      betas (Tuple[float, float], optional): coefficients used for computing
          running averages of gradient and its square (default: (0.9, 0.999))
      eps (float, optional): term added to the denominator to improve
          numerical stability (default: 1e-8)
      weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
      amsgrad (boolean, optional): whether to use the AMSGrad variant of this
          algorithm from the paper `On the Convergence of Adam and Beyond`_
  
  .. _Adam\: A Method for Stochastic Optimization:
      https://arxiv.org/abs/1412.6980
  .. _On the Convergence of Adam and Beyond:
      https://openreview.net/forum?id=ryQu7f-RZ
class_docstr_tok: ["Implements","Adam","algorithm",".","This","implementation","is","modified","from","torch.optim.Adam","based","on",":","`","Fixed","Weight","Decay","Regularization","in","Adam","`","(","see","https",":","\/\/arxiv.org\/abs\/1711.05101",")","It","has","been","proposed","in","`","Adam",":","A","Method","for","Stochastic","Optimization","`","_",".","Args",":","params","(","iterable",")",":","iterable","of","parameters","to","optimize","or","dicts","defining","parameter","groups","lr","(","float",",","optional",")",":","learning","rate","(","default",":","1e-3",")","betas","(","Tuple","[","float",",","float","]",",","optional",")",":","coefficients","used","for","computing","running","averages","of","gradient","and","its","square","(","default",":","(","0.9",",","0.999",")",")","eps","(","float",",","optional",")",":","term","added","to","the","denominator","to","improve","numerical","stability","(","default",":","1e-8",")","weight_decay","(","float",",","optional",")",":","weight","decay","(","L2","penalty",")","(","default",":","0",")","amsgrad","(","boolean",",","optional",")",":","whether","to","use","the","AMSGrad","variant","of","this","algorithm","from","the","paper","`","On","the","Convergence","of","Adam","and","Beyond","`","_","..","_Adam\\",":","A","Method","for","Stochastic","Optimization",":","https",":","\/\/arxiv.org\/abs\/1412.6980","..","_On","the","Convergence","of","Adam","and","Beyond",":","https",":","\/\/openreview.net\/forum","?","id=ryQu7f-RZ"]
docstr: Performs a single optimization step.
  
  Arguments:
      closure (callable, optional): A closure that reevaluates the model
          and returns the loss.
docstr_tok: ["Performs","a","single","optimization","step",".","Arguments",":","closure","(","callable",",","optional",")",":","A","closure","that","reevaluates","the","model","and","returns","the","loss","."]
code: def step(self, closure=None):
      """Performs a single optimization step.
  
      Arguments:
          closure (callable, optional): A closure that reevaluates the model
              and returns the loss.
      """
      loss = None
      if closure is not None:
          loss = closure()
  
      for group in self.param_groups:
          for p in group["params"]:
              if p.grad is None:
                  continue
              grad = p.grad.data
              if grad.dtype in {torch.float16, torch.bfloat16}:
                  grad = grad.float()
              if grad.is_sparse:
                  raise RuntimeError(
                      "Adam does not support sparse gradients, please consider SparseAdam instead"
                  )
              amsgrad = group.get("amsgrad", False)
  
              p_data_fp32 = p.data
              if p.data.dtype in {torch.float16, torch.bfloat16}:
                  p_data_fp32 = p_data_fp32.float()
  
              state = self.state[p]
  
              # State initialization
              if len(state) == 0:
                  state["step"] = 0
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(p_data_fp32)
                  # Exponential moving average of squared gradient values
                  state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
                  if amsgrad:
                      # Maintains max of all exp. moving avg. of sq. grad. values
                      state["max_exp_avg_sq"] = torch.zeros_like(
                          p_data_fp32
                      )
              else:
                  state["exp_avg"] = state["exp_avg"].to(p_data_fp32)
                  state["exp_avg_sq"] = state["exp_avg_sq"].to(
                      p_data_fp32
                  )
                  if amsgrad:
                      state["max_exp_avg_sq"] = state[
                          "max_exp_avg_sq"
                      ].to(p_data_fp32)
  
              exp_avg, exp_avg_sq = (
                  state["exp_avg"],
                  state["exp_avg_sq"],
              )
              if amsgrad:
                  max_exp_avg_sq = state["max_exp_avg_sq"]
              beta1, beta2 = group["betas"]
  
              state["step"] += 1
  
              # Decay the first and second moment running average coefficient
              exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
              exp_avg_sq.mul_(beta2).addcmul_(
                  grad, grad, value=1 - beta2
              )
              if amsgrad:
                  # Maintains the maximum of all 2nd moment running avg. till now
                  torch.max(
                      max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq
                  )
                  # Use the max. for normalizing running avg. of gradient
                  denom = max_exp_avg_sq.sqrt().add_(group["eps"])
              else:
                  denom = exp_avg_sq.sqrt().add_(group["eps"])
  
              bias_correction1 = 1 - beta1 ** state["step"]
              bias_correction2 = 1 - beta2 ** state["step"]
              step_size = (
                  group["lr"]
                  * math.sqrt(bias_correction2)
                  / bias_correction1
              )
  
              if group["weight_decay"] != 0:
                  p_data_fp32.add_(
                      p_data_fp32,
                      alpha=-group["weight_decay"] * group["lr"],
                  )
  
              p_data_fp32.addcdiv_(exp_avg, denom, value=-step_size)
  
              if p.data.dtype in {torch.float16, torch.bfloat16}:
                  p.data.copy_(p_data_fp32)
  
      return loss
  
imports: import logging
  import math
  from dataclasses import dataclass, field
  from typing import List
  import torch
  import torch.distributed
  import torch.optim
  from fairseq.dataclass import FairseqDataclass
  from fairseq.optim import FairseqOptimizer, register_optimizer
  from omegaconf import II
  from lightseq.training.ops.pytorch.adam import LSAdam
label_desc: Use image node_id 4 for calling the Adam obj's underlying member method code with example usage: obj.step(closure) and returns: loss