class_docstr: Implements Adafactor algorithm.
  
  This implementation is based on:
  `Adafactor: Adaptive Learning Rates with Sublinear Memory Cost`
  (see https://arxiv.org/abs/1804.04235)
  
  Taken from the fairseq implementation, which can be found here:
  <https://github.com/pytorch/fairseq/blob/main/fairseq/optim/adafactor.py>.
  
  :param params (iterable):
      iterable of parameters to optimize or dicts defining parameter groups
  :param lr (float, optional):
      external learning rate (default: None)
  :param eps (tuple[float, float]):
      regularization constans for square gradient and parameter scale
      respectively (default: (1e-30, 1e-3))
  :param clip_threshold (float):
      threshold of root mean square of final gradient update
      (default: 1.0)
  :param decay_rate (float):
      coefficient used to compute running averages of square gradient
      (default: -0.8)
  :param beta1 (float):
      coefficient used for computing running averages of gradient
      (default: None)
  :param weight_decay (float, optional):
      weight decay (L2 penalty) (default: 0)
  :param scale_parameter (bool):
      if true, learning rate is scaled by root mean square of parameter
      (default: True)
  :param relative_step (bool):
      if true, time-dependent learning rate is computed instead of external
      learning rate (default: True)
  :param warmup_init (bool):
      time-dependent learning rate computation depends on whether warm-up
      initialization is being used (default: False)
class_docstr_tok: ["Implements","Adafactor","algorithm",".","This","implementation","is","based","on",":","`","Adafactor",":","Adaptive","Learning","Rates","with","Sublinear","Memory","Cost","`","(","see","https",":","\/\/arxiv.org\/abs\/1804.04235",")","Taken","from","the","fairseq","implementation",",","which","can","be","found","here",":","<","https",":","\/\/github.com\/pytorch\/fairseq\/blob\/main\/fairseq\/optim\/adafactor.py",">",".",":","param","params","(","iterable",")",":","iterable","of","parameters","to","optimize","or","dicts","defining","parameter","groups",":","param","lr","(","float",",","optional",")",":","external","learning","rate","(","default",":","None",")",":","param","eps","(","tuple","[","float",",","float","]",")",":","regularization","constans","for","square","gradient","and","parameter","scale","respectively","(","default",":","(","1e-30",",","1e-3",")",")",":","param","clip_threshold","(","float",")",":","threshold","of","root","mean","square","of","final","gradient","update","(","default",":","1.0",")",":","param","decay_rate","(","float",")",":","coefficient","used","to","compute","running","averages","of","square","gradient","(","default",":","-0.8",")",":","param","beta1","(","float",")",":","coefficient","used","for","computing","running","averages","of","gradient","(","default",":","None",")",":","param","weight_decay","(","float",",","optional",")",":","weight","decay","(","L2","penalty",")","(","default",":","0",")",":","param","scale_parameter","(","bool",")",":","if","true",",","learning","rate","is","scaled","by","root","mean","square","of","parameter","(","default",":","True",")",":","param","relative_step","(","bool",")",":","if","true",",","time-dependent","learning","rate","is","computed","instead","of","external","learning","rate","(","default",":","True",")",":","param","warmup_init","(","bool",")",":","time-dependent","learning","rate","computation","depends","on","whether","warm-up","initialization","is","being","used","(","default",":","False",")"]
docstr: 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.float()  # NOTE: cast to FP32
              if grad.is_sparse:
                  raise RuntimeError(
                      "Adafactor does not support sparse gradients."
                  )
  
              state = self.state[p]
              grad_shape = grad.shape
  
              factored, use_first_moment = self._get_options(
                  group, grad_shape
              )
              # State Initialization
              if len(state) == 0:
                  state["step"] = 0
  
                  if use_first_moment:
                      # Exponential moving average of gradient values
                      state["exp_avg"] = torch.zeros_like(grad)
                  if factored:
                      state["exp_avg_sq_row"] = torch.zeros(
                          grad_shape[:-1]
                      ).type_as(grad)
                      state["exp_avg_sq_col"] = torch.zeros(
                          grad_shape[:-2] + grad_shape[-1:]
                      ).type_as(grad)
                  else:
                      state["exp_avg_sq"] = torch.zeros_like(grad)
  
                  state["RMS"] = 0
              else:
                  if use_first_moment:
                      state["exp_avg"] = state["exp_avg"].type_as(grad)
                  if factored:
                      state["exp_avg_sq_row"] = state[
                          "exp_avg_sq_row"
                      ].type_as(grad)
                      state["exp_avg_sq_col"] = state[
                          "exp_avg_sq_col"
                      ].type_as(grad)
                  else:
                      state["exp_avg_sq"] = state["exp_avg_sq"].type_as(
                          grad
                      )
  
              p_data_fp32 = p.data.float()  # NOTE: cast to FP32
  
              state["step"] += 1
              state["RMS"] = self._rms(p_data_fp32)
              group["lr"] = self._get_lr(group, state)
  
              beta2t = 1.0 - math.pow(
                  state["step"], group["decay_rate"]
              )
              update = (grad**2) + group["eps"][0]
              if factored:
                  exp_avg_sq_row = state["exp_avg_sq_row"]
                  exp_avg_sq_col = state["exp_avg_sq_col"]
  
                  exp_avg_sq_row.mul_(beta2t).add_(
                      update.mean(dim=-1), alpha=(1.0 - beta2t)
                  )
                  exp_avg_sq_col.mul_(beta2t).add_(
                      update.mean(dim=-2), alpha=(1.0 - beta2t)
                  )
  
                  # Approximation of exponential moving average of square of gradient
                  self._approx_sq_grad(
                      exp_avg_sq_row, exp_avg_sq_col, update
                  )
                  update.mul_(grad)
              else:
                  exp_avg_sq = state["exp_avg_sq"]
  
                  exp_avg_sq.mul_(beta2t).add_(
                      update, alpha=(1.0 - beta2t)
                  )
                  torch.rsqrt(exp_avg_sq, out=update).mul_(grad)
  
              update.div_(
                  max(1.0, self._rms(update) / group["clip_threshold"])
              )
              update.mul_(group["lr"])
  
              if use_first_moment:
                  exp_avg = state["exp_avg"]
                  exp_avg.mul_(group["beta1"]).add_(
                      update, alpha=(1 - group["beta1"])
                  )
                  update = exp_avg
  
              if group["weight_decay"] != 0:
                  p_data_fp32.add_(
                      p_data_fp32,
                      alpha=(-group["weight_decay"] * group["lr"]),
                  )
  
              p_data_fp32.add_(-update)
  
              p.data.copy_(p_data_fp32)
  
      return loss
  
imports: import math
  from itertools import chain
  from typing import Optional
  import parlai.utils.logging
label_desc: Use image node_id 6 for calling the Adafactor obj's underlying member method code with example usage: obj.step(closure) and returns: loss