class_docstr: Implements Adafactor algorithm.
  
  It has been proposed in: `Adafactor: Adaptive Learning Rates with
  Sublinear Memory Cost`__.
  
  Arguments:
      params: iterable of parameters to optimize or dicts defining
          parameter groups
      lr: external learning rate (default: None)
      eps2: regularization constans for square gradient
          and parameter scale respectively (default: (1e-30, 1e-3))
      clip_threshold: threshold of root mean square of
          final gradient update (default: 1.0)
      decay_rate: coefficient used to compute running averages of square
          gradient (default: -0.8)
      beta1: coefficient used for computing running averages of gradient
          (default: None)
      weight_decay: weight decay (L2 penalty) (default: 0)
      scale_parameter: if true, learning rate is scaled by root mean square
          of parameter (default: True)
      relative_step: if true, time-dependent learning rate is computed
          instead of external learning rate (default: True)
      warmup_init: time-dependent learning rate computation depends on
          whether warm-up initialization is being used (default: False)
  
  Example:
      >>> import torch_optimizer as optim
      >>> optimizer = optim.Adafactor(model.parameters())
      >>> optimizer.zero_grad()
      >>> loss_fn(model(input), target).backward()
      >>> optimizer.step()
  
  __ https://arxiv.org/abs/1804.04235
  
  Note:
      Reference code: https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py  # noqa
class_docstr_tok: ["Implements","Adafactor","algorithm",".","It","has","been","proposed","in",":","`","Adafactor",":","Adaptive","Learning","Rates","with","Sublinear","Memory","Cost","`","__",".","Arguments",":","params",":","iterable","of","parameters","to","optimize","or","dicts","defining","parameter","groups","lr",":","external","learning","rate","(","default",":","None",")","eps2",":","regularization","constans","for","square","gradient","and","parameter","scale","respectively","(","default",":","(","1e-30",",","1e-3",")",")","clip_threshold",":","threshold","of","root","mean","square","of","final","gradient","update","(","default",":","1.0",")","decay_rate",":","coefficient","used","to","compute","running","averages","of","square","gradient","(","default",":","-0.8",")","beta1",":","coefficient","used","for","computing","running","averages","of","gradient","(","default",":","None",")","weight_decay",":","weight","decay","(","L2","penalty",")","(","default",":","0",")","scale_parameter",":","if","true",",","learning","rate","is","scaled","by","root","mean","square","of","parameter","(","default",":","True",")","relative_step",":","if","true",",","time-dependent","learning","rate","is","computed","instead","of","external","learning","rate","(","default",":","True",")","warmup_init",":","time-dependent","learning","rate","computation","depends","on","whether","warm-up","initialization","is","being","used","(","default",":","False",")","Example",":",">",">",">","import","torch_optimizer","as","optim",">",">",">","optimizer","=","optim.Adafactor","(","model.parameters","(",")",")",">",">",">","optimizer.zero_grad","(",")",">",">",">","loss_fn","(","model","(","input",")",",","target",")",".backward","(",")",">",">",">","optimizer.step","(",")","__","https",":","\/\/arxiv.org\/abs\/1804.04235","Note",":","Reference","code",":","https",":","\/\/github.com\/pytorch\/fairseq\/blob\/master\/fairseq\/optim\/adafactor.py","#","noqa"]
docstr: Performs a single optimization step.
  
  Arguments:
      closure: A closure that reevaluates the model and returns the loss.
docstr_tok: ["Performs","a","single","optimization","step",".","Arguments",":","closure",":","A","closure","that","reevaluates","the","model","and","returns","the","loss","."]
code: def step(self, closure: OptLossClosure = None) -> OptFloat:
      r"""Performs a single optimization step.
  
      Arguments:
          closure: 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.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, memory_format=torch.preserve_format
                      )
                  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, memory_format=torch.preserve_format
                      )
  
                  state["RMS"] = 0
  
              state["step"] += 1
              state["RMS"] = self._rms(p.data)
              lr = self._get_lr(group, state)
  
              beta2t = 1.0 - math.pow(
                  state["step"], group["decay_rate"]
              )
              update = (grad**2) + group["eps2"][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_(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.add_(p.data, alpha=-group["weight_decay"] * lr)
  
              p.data.add_(-update)
  
      return loss
  
file_path: pytorch-optimizer/torch_optimizer/adafactor.py
imports: import math
  from typing import Any, Dict, Tuple
  import torch
  from torch.optim.optimizer import Optimizer
  from .types import OptFloat, OptLossClosure, Params, State
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