code: def _build_pointer_buffers(self):
      # Loading checkpoint to optimizer re-allocates param and states, so pointer logic should be
      # at the first step of training, where we assume all states are ready.
      if not all(
          len(pg) == 1
          for pg in (
              self.param_groups,
              self.compute_param_groups,
              self.swa_param_groups,
          )
      ):
          raise RuntimeError(
              "FusedAdamSWA does not support multiple param groups"
          )
  
      # `bf16_params` contains both BF16 and FP32 data types, thus we have to group parameters
      # and other states into different buffers and launch respective kernels.
      params, compute_params, swa_params = (
          self.param_groups[0]["params"],
          self.compute_param_groups[0]["params"],
          self.swa_param_groups[0]["params"],
      )
      self.pointer_buffer_groups = defaultdict(dict)
      for i, p in enumerate(compute_params):
          compute_dtype = p.dtype
          state_dtype = params[i].dtype
          self.pointer_buffer_groups[
              (compute_dtype, state_dtype)
          ].setdefault("tensor_idx", [])
          self.pointer_buffer_groups[(compute_dtype, state_dtype)][
              "tensor_idx"
          ].append(i)
  
      for (
          _,
          state_dtype,
      ), buffer_group in self.pointer_buffer_groups.items():
          # Select tensors by dtype.
          t_idx = buffer_group["tensor_idx"]
          params_this_group = [params[i] for i in t_idx]
          compute_params_this_group = [compute_params[i] for i in t_idx]
          swa_params_this_group = [swa_params[i] for i in t_idx]
  
          # Build parameter pointer buffers.
          param_ptrs = torch.tensor(
              [p.data_ptr() for p in params_this_group],
              dtype=torch.int64,
          )
          compute_param_ptrs = torch.tensor(
              [b.data_ptr() for b in compute_params_this_group],
              dtype=torch.int64,
          )
          swa_param_ptrs = torch.tensor(
              [s.data_ptr() for s in swa_params_this_group],
              dtype=torch.int64,
          )
  
          param_numels = torch.tensor(
              [p.numel() for p in params_this_group], dtype=torch.int64
          )
          chunks_per_param = (
              param_numels.float().div_(CHUNK_SIZE).ceil_().long()
          )
          chunk_local_idx = torch.cat(
              [
                  torch.arange(chunks, dtype=torch.int64)
                  for chunks in chunks_per_param
              ]
          )
          chunk_numel = torch.minimum(
              param_numels.repeat_interleave(chunks_per_param)
              - chunk_local_idx * CHUNK_SIZE,
              CHUNK_SIZE,
          )
          param_ptr_per_chunk = torch.repeat_interleave(
              param_ptrs, chunks_per_param
          )
          compute_param_ptr_per_chunk = torch.repeat_interleave(
              compute_param_ptrs, chunks_per_param
          )
          swa_param_ptr_per_chunk = torch.repeat_interleave(
              swa_param_ptrs, chunks_per_param
          )
  
          device = params_this_group[0].device
          buffer_group["device"] = device
          buffer_group["chunks_per_param"] = chunks_per_param
          buffer_group["chunk_local_idx"] = chunk_local_idx.to(device)
          buffer_group["chunk_numel"] = chunk_numel.to(device)
          buffer_group["param_ptr_per_chunk"] = param_ptr_per_chunk.to(
              device
          )
          buffer_group[
              "compute_param_ptr_per_chunk"
          ] = compute_param_ptr_per_chunk.to(device)
          buffer_group[
              "swa_param_ptr_per_chunk"
          ] = swa_param_ptr_per_chunk.to(device)
          buffer_group["total_chunks"] = chunks_per_param.sum().item()
          buffer_group["default_grad_clip_scale"] = torch.tensor(
              1.0, dtype=state_dtype
          ).to(device)
  
          # Build moment pointer buffers.
          moment, velocity = [], []
          for p in params_this_group:
              state = self.state[p]
              if "exp_avg" not in state or "exp_avg_sq" not in state:
                  state["exp_avg"] = torch.zeros_like(
                      p.detach(), dtype=state_dtype
                  )
                  state["exp_avg_sq"] = torch.zeros_like(
                      p.detach(), dtype=state_dtype
                  )
              moment.append(state["exp_avg"].data_ptr())
              velocity.append(state["exp_avg_sq"].data_ptr())
          moment = torch.tensor(moment, dtype=torch.int64)
          velocity = torch.tensor(velocity, dtype=torch.int64)
          buffer_group[
              "exp_avg_ptr_per_chunk"
          ] = torch.repeat_interleave(moment, chunks_per_param).to(
              device
          )
          buffer_group[
              "exp_avg_sq_ptr_per_chunk"
          ] = torch.repeat_interleave(velocity, chunks_per_param).to(
              device
          )
  
      self._pointer_buffers_initialized = True
  
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 2 for calling the FusedAdamSWA obj's underlying member method code with example usage: obj._build_pointer_buffers() without return types