class_docstr: Implements AdamW algorithm.
  
  This class uses the same functional AdamW implementation as
  :class:`torch.optim.AdamW`. The main difference is that it keeps its state
  in single precision (i.e. ``torch.float32``), and temporarily casts half
  precision parameters and gradients (i.e. ``torch.float16`` and
  ``torch.bfloat16``) to single precision during an optimization step for
  better numerical stability.
class_docstr_tok: ["Implements","AdamW","algorithm",".","This","class","uses","the","same","functional","AdamW","implementation","as",":","class",":","`","torch.optim.AdamW","`",".","The","main","difference","is","that","it","keeps","its","state","in","single","precision","(","i.e",".","``","torch.float32","``",")",",","and","temporarily","casts","half","precision","parameters","and","gradients","(","i.e",".","``","torch.float16","``","and","``","torch.bfloat16","``",")","to","single","precision","during","an","optimization","step","for","better","numerical","stability","."]
docstr: :param params:
      The parameters to optimize.
  :param lr:
      The learning rate.
  :param betas:
      The coefficients used for computing running averages of gradient and
      its square.
  :param eps:
      The term added to the denominator to improve numerical stability.
  :param weight_decay:
      The weight decay coefficient.
  :param amsgrad:
      If ``True``, uses the AMSGrad variant.
  :param maximize:
      If ``True``, maximizes the parameters based on the objective,
      instead of minimizing.
  :param capturable:
      If ``True``, indicates that it is safe to capture this instance in a
      CUDA graph.
  :param differentiable:
      If ``True``, runs the optimizer step under autograd.
  :param impl:
      The implementation variant. See :class:`torch.optim.AdamW` for the
      details.
docstr_tok: [":","param","params",":","The","parameters","to","optimize",".",":","param","lr",":","The","learning","rate",".",":","param","betas",":","The","coefficients","used","for","computing","running","averages","of","gradient","and","its","square",".",":","param","eps",":","The","term","added","to","the","denominator","to","improve","numerical","stability",".",":","param","weight_decay",":","The","weight","decay","coefficient",".",":","param","amsgrad",":","If","``","True","``",",","uses","the","AMSGrad","variant",".",":","param","maximize",":","If","``","True","``",",","maximizes","the","parameters","based","on","the","objective",",","instead","of","minimizing",".",":","param","capturable",":","If","``","True","``",",","indicates","that","it","is","safe","to","capture","this","instance","in","a","CUDA","graph",".",":","param","differentiable",":","If","``","True","``",",","runs","the","optimizer","step","under","autograd",".",":","param","impl",":","The","implementation","variant",".","See",":","class",":","`","torch.optim.AdamW","`","for","the","details","."]
code: def __init__(
      self,
      params: Union[Iterable[Tensor], Iterable[Dict[str, Any]]],
      *,
      lr: float = 1e-3,
      betas: Tuple[float, float] = (0.9, 0.999),
      eps: float = 1e-8,
      weight_decay: float = 0.0,
      amsgrad: bool = False,
      maximize: bool = False,
      capturable: bool = False,
      differentiable: bool = False,
      impl: Literal["auto", "foreach", "fused"] = "auto",
  ) -> None:
      """
      :param params:
          The parameters to optimize.
      :param lr:
          The learning rate.
      :param betas:
          The coefficients used for computing running averages of gradient and
          its square.
      :param eps:
          The term added to the denominator to improve numerical stability.
      :param weight_decay:
          The weight decay coefficient.
      :param amsgrad:
          If ``True``, uses the AMSGrad variant.
      :param maximize:
          If ``True``, maximizes the parameters based on the objective,
          instead of minimizing.
      :param capturable:
          If ``True``, indicates that it is safe to capture this instance in a
          CUDA graph.
      :param differentiable:
          If ``True``, runs the optimizer step under autograd.
      :param impl:
          The implementation variant. See :class:`torch.optim.AdamW` for the
          details.
      """
      defaults = {
          "lr": lr,
          "betas": betas,
          "eps": eps,
          "weight_decay": weight_decay,
          "amsgrad": amsgrad,
          "maximize": maximize,
          "capturable": capturable,
          "differentiable": differentiable,
          "impl": impl,
      }
  
      super().__init__(params, defaults)
  
      if impl == "fused":
          if differentiable:
              raise RuntimeError(
                  "`fused` implementation does not support `differentiable`."
              )
  
          for pg in self.param_groups:
              for p in pg["params"]:
                  if (
                      not torch.is_floating_point(p)
                      or p.device.type != "cuda"
                  ):
                      raise RuntimeError(
                          "`fused` implementation requires all parameters to be float CUDA tensors."
                      )
  
          self._step_supports_amp_scaling = True
  
file_path: fairseq2/src/fairseq2/optim/adamw.py
imports: from itertools import chain
  from typing import Any, Dict, Iterable, List, Literal, Tuple, Union, cast, final
  import torch
  from torch import Tensor
  from torch.optim.adamw import adamw
  from fairseq2.optim.optimizer_base import OptimizerBase
  from fairseq2.typing import finaloverride
label_desc: Use image node_id 1 to create a new AdamW object from inherited base classes: OptimizerBase with example: obj = AdamW(params)