code: def __init__(
      self,
      params: List[nn.Parameter],
      compute_params: List[nn.Parameter],
      swa_params: List[nn.Parameter],
      swa_decay_rate: float,
      lr: float = 1e-3,
      bias_correction: bool = True,
      betas: Tuple[float, float] = (0.9, 0.999),
      eps: float = 1e-8,
      adam_math_mode: AdamMathType = AdamMathType.PyTorchAdam,
      weight_decay: float = 0.0,
      amsgrad: bool = False,
      set_grad_none: bool = True,
      capturable: bool = False,
      master_weights: bool = False,
  ):
      if not isinstance(params, list):
          params = list(params)
      if not isinstance(compute_params, list):
          compute_params = list(compute_params)
      if not isinstance(swa_params, list):
          swa_params = list(swa_params)
      if not compute_params or not swa_params:
          raise ValueError(
              "FusedAdamSWA requires both BF16 and SWA parameters."
          )
      if not len(params) == len(compute_params) == len(swa_params):
          raise ValueError(
              "FusedAdamSWA expects params, bf16_params, and swa_params to have same length"
          )
      if not all(
          p.shape == b.shape == s.shape
          for p, b, s in zip(params, compute_params, swa_params)
      ):
          raise ValueError(
              "FusedAdamSWA expects each state in params, bf16_params, abd swa_params to have same shape"
          )
      if not all(
          p.dtype == s.dtype for p, s in zip(params, swa_params)
      ):
          raise ValueError(
              "FusedAdamSWA expects all params and swa_params to have same dtype"
          )
      if not all(
          p.is_contiguous()
          for p in chain(params, compute_params, swa_params)
      ):
          raise ValueError(
              "FusedAdamSWA expects all input params to be contiguous"
          )
      if amsgrad:
          raise NotImplementedError(
              "amsgrad is not supported by FusedAdamSWA"
          )
      if capturable:
          raise NotImplementedError(
              "capturable is not supported by FusedAdamSWA"
          )
      if master_weights:
          raise NotImplementedError(
              "master_weights is not supported by FusedAdamSWA"
          )
      if not isinstance(adam_math_mode, AdamMathType):
          raise ValueError(
              f"Unknown Adam math mode {adam_math_mode}, expect to be any of:\n"
              f"\t- {AdamMathType.ApexAdam}: NVIDIA Apex Adam math;\n"
              f"\t- {AdamMathType.ApexAdamW}: NVIDIA Apex Adam math with adam_w set to True;\n"
              f"\t- {AdamMathType.PyTorchAdam}: The official PyTorch Adam math.\n"
          )
  
      defaults = dict(
          lr=lr,
          bias_correction=bias_correction,
          betas=betas,
          eps=eps,
          weight_decay=weight_decay,
      )
      super().__init__(params, defaults)
      self.adam_math_mode = adam_math_mode
      self.set_grad_none = set_grad_none
      self.compute_param_groups = [{"params": compute_params}]
      self.swa_param_groups = [{"params": swa_params, "n_averaged": 0}]
      self.swa_decay_rate = swa_decay_rate
  
      # We assume that parameter and buffer pointers won't change throughout the training, only
      # gradients could be re-allocated due to set_grad_none.
      self._pointer_buffers_initialized = False
  
imports: from __future__ import annotations
  from collections import defaultdict
  from enum import Enum, unique
  from itertools import chain
  from typing import Callable, List, Optional, Tuple, Union
  import torch
  import torch.nn
  import triton
  import triton.language
  from torch.optim import Adam, Optimizer
label_desc: Use image node_id 1 to create a new FusedAdamSWA object from inherited base classes: Optimizer with example: obj = FusedAdamSWA(params, compute_params, swa_params, swa_decay_rate, lr, bias_correction, betas, eps, adam_math_mode, weight_decay, amsgrad, set_grad_none, capturable, master_weights)