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","."]
code: def _init_param(
      self,
      param: Tensor,
      param_group: Dict[str, Any],
      params_with_grad: List[Tensor],
      fp32_params_with_grad: List[Tensor],
      fp32_grads: List[Tensor],
      steps: List[Tensor],
      exp_avgs: List[Tensor],
      exp_avg_sqs: List[Tensor],
      max_exp_avg_sqs: List[Tensor],
      amsgrad: bool,
  ) -> None:
      grad = param.grad
      if grad is None:
          return
  
      if grad.is_sparse:
          raise RuntimeError(
              "`AdamW` does not support sparse gradients."
          )
  
      params_with_grad.append(param)
  
      # fp32 parameter
      if param.dtype == torch.float16 or param.dtype == torch.bfloat16:
          fp32_param = param.float()
      else:
          fp32_param = param
  
      fp32_params_with_grad.append(fp32_param)
  
      # fp32 grad
      if grad.dtype == torch.float16 or grad.dtype == torch.bfloat16:
          fp32_grads.append(grad.float())
      else:
          fp32_grads.append(grad)
  
      state = cast(Dict[str, Tensor], self.state[param])  # type: ignore[index]
  
      if len(state) == 0:
          if (
              param_group["capturable"]
              or param_group["impl"] == "fused"
          ):
              step_device = param.device
          else:
              step_device = None
  
          # Step counter.
          state["step"] = torch.zeros(
              (), device=step_device, dtype=torch.float32
          )
  
          # Exponential moving average of gradient values.
          state["exp_avg"] = torch.zeros_like(fp32_param)
  
          # Exponential moving average of squared gradient values.
          state["exp_avg_sq"] = torch.zeros_like(fp32_param)
  
          if amsgrad:
              state["max_exp_avg_sq"] = torch.zeros_like(fp32_param)
  
      steps.append(state["step"])
  
      exp_avgs.append(state["exp_avg"])
  
      exp_avg_sqs.append(state["exp_avg_sq"])
  
      if amsgrad:
          max_exp_avg_sqs.append(state["max_exp_avg_sq"])
  
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 4 for calling the AdamW obj's underlying member method code with example usage: obj._init_param(param, param_group, params_with_grad, fp32_params_with_grad, fp32_grads, steps, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, amsgrad) without return types