code: def _multi_tensor_adam_swa(
      state_param_ptr_per_chunk,
      compute_param_ptr_per_chunk,
      swa_param_ptr_per_chunk,
      grad_ptr_per_chunk,
      moment_ptr_per_chunk,
      velocity_ptr_per_chunk,
      chunk_local_idx_ptr,
      chunk_numel_ptr,
      grad_clip_scale_ptr,
      lr,
      beta1,
      beta2,
      eps,
      weight_decay,
      beta1_correction,
      beta2_correction,
      swa_decay_rate,
      swa_n_averaged,
      adam_math_mode: tl.constexpr,
      MODEL_COMPUTE_DTYPE: tl.constexpr,
      MODEL_STATE_DTYPE: tl.constexpr,
      CHUNK_SIZE: tl.constexpr,
      BLOCK_SIZE: tl.constexpr,
  ):
      chunk_idx = tl.program_id(0)
      chunk_local_idx = tl.load(chunk_local_idx_ptr + chunk_idx)
      chunk_numel = tl.load(chunk_numel_ptr + chunk_idx)
  
      compute_dtype = _DTYPE2TRITON[MODEL_COMPUTE_DTYPE.value]
      compute_pointer_type = tl.pointer_type(compute_dtype)
      state_dtype = _DTYPE2TRITON[MODEL_STATE_DTYPE.value]
      state_pointer_type = tl.pointer_type(state_dtype)
  
      state_param_ptr = tl.load(
          state_param_ptr_per_chunk + chunk_idx
      ).to(state_pointer_type)
      swa_param_ptr = tl.load(swa_param_ptr_per_chunk + chunk_idx).to(
          state_pointer_type
      )
      moment_ptr = tl.load(moment_ptr_per_chunk + chunk_idx).to(
          state_pointer_type
      )
      velocity_ptr = tl.load(velocity_ptr_per_chunk + chunk_idx).to(
          state_pointer_type
      )
      compute_param_ptr = tl.load(
          compute_param_ptr_per_chunk + chunk_idx
      ).to(compute_pointer_type)
      grad_ptr = tl.load(grad_ptr_per_chunk + chunk_idx).to(
          compute_pointer_type
      )
      grad_clip_scale = tl.load(grad_clip_scale_ptr)
  
      ptr_base_offset = chunk_local_idx * CHUNK_SIZE
      state_param_ptr += ptr_base_offset
      compute_param_ptr += ptr_base_offset
      swa_param_ptr += ptr_base_offset
      grad_ptr += ptr_base_offset
      moment_ptr += ptr_base_offset
      velocity_ptr += ptr_base_offset
  
      for i in range(0, CHUNK_SIZE, BLOCK_SIZE):
          idx = i + tl.arange(0, BLOCK_SIZE)
          mask = idx < chunk_numel
          # Gradient clip step.
          grad = tl.load(grad_ptr + idx, mask).to(state_dtype)
          grad *= grad_clip_scale
          # Adam step.
          param = tl.load(state_param_ptr + idx, mask)
          moment = tl.load(moment_ptr + idx, mask)
          velocity = tl.load(velocity_ptr + idx, mask)
          param, moment, velocity = _adam_math(
              param=param,
              grad=grad,
              moment=moment,
              velocity=velocity,
              beta1=beta1,
              beta2=beta2,
              beta1_correction=beta1_correction,
              beta2_correction=beta2_correction,
              eps=eps,
              lr=lr,
              weight_decay=weight_decay,
              adam_math_mode=adam_math_mode,
          )
          # SWA step.
          swa_param = tl.load(swa_param_ptr + idx, mask)
          swa_param = _swa_math(
              param=param,
              swa_param=swa_param,
              decay_rate=swa_decay_rate,
              n_averaged=swa_n_averaged,
          )
          # Write results. BF16 and SWA parameters are updated as well.
          tl.store(state_param_ptr + idx, param, mask)
          tl.store(moment_ptr + idx, moment, mask)
          tl.store(velocity_ptr + idx, velocity, mask)
          tl.store(compute_param_ptr + idx, param, mask)
          tl.store(swa_param_ptr + idx, swa_param, mask)
  
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 3 for calling a global function with example usage: _multi_tensor_adam_swa(state_param_ptr_per_chunk, compute_param_ptr_per_chunk, swa_param_ptr_per_chunk, grad_ptr_per_chunk, moment_ptr_per_chunk, velocity_ptr_per_chunk, chunk_local_idx_ptr, chunk_numel_ptr, grad_clip_scale_ptr, lr, beta1, beta2, eps, weight_decay, beta1_correction, beta2_correction, swa_decay_rate, swa_n_averaged, adam_math_mode, MODEL_COMPUTE_DTYPE, MODEL_STATE_DTYPE, CHUNK_SIZE, BLOCK_SIZE) without return types