code: def from_optim(
      cls,
      adam_optimizer: Adam,
      fp32_params: List[nn.Parameter],
      bf16_params: List[nn.Parameter],
      swa_params: List[nn.Parameter],
      swa_decay_rate: float,
  ) -> FusedAdamSWA:
      assert len(adam_optimizer.param_groups) == 1
      param_group = adam_optimizer.param_groups[0]
      lr = param_group["lr"]
      betas = param_group["betas"]
      eps = param_group["eps"]
      weight_decay = param_group["weight_decay"]
      amsgrad = param_group["amsgrad"]
      fused_adam_swa_optimizer = cls(
          params=fp32_params,
          compute_params=bf16_params,
          swa_params=swa_params,
          swa_decay_rate=swa_decay_rate,
          lr=lr,
          betas=betas,
          eps=eps,
          weight_decay=weight_decay,
          amsgrad=amsgrad,
          adam_math_mode=AdamMathType.PyTorchAdam,
      )
      adam_state_dict = adam_optimizer.state_dict()
      adam_state_dict["param_groups"][0].setdefault(
          "bias_correction", True
      )
      steps = [v["step"] for v in adam_state_dict["state"].values()]
      if len(steps) == 0:  # Did not load optimizer checkpoint.
          steps = [torch.tensor(1)]
      elif not all(s == steps[0] for s in steps):
          raise ValueError(
              "FusedAdamSWA requires all parameters were updated by same steps!"
          )
      step = int(steps[0].item())
      adam_state_dict["param_groups"][0].setdefault("step", step)
      fused_adam_swa_optimizer.load_state_dict(adam_state_dict)
      return fused_adam_swa_optimizer
  
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 4 for calling the FusedAdamSWA obj's underlying member method code with example usage: obj.from_optim(cls, adam_optimizer, fp32_params, bf16_params, swa_params, swa_decay_rate) and returns: fused_adam_swa_optimizer