code: def _adam_math(
      param,
      grad,
      moment,
      velocity,
      beta1,
      beta2,
      beta1_correction,
      beta2_correction,
      eps,
      lr,
      weight_decay,
      adam_math_mode: tl.constexpr,
  ):
      if adam_math_mode == tl.constexpr(AdamMathType.ApexAdam.value):
          grad += weight_decay * param
          moment *= beta1
          moment += (1.0 - beta1) * grad
          velocity *= beta2
          velocity += (1.0 - beta2) * grad * grad
          update = (moment / beta1_correction) / (
              tl.math.sqrt(velocity / beta2_correction) + eps
          )
          param -= lr * update
      elif adam_math_mode == tl.constexpr(AdamMathType.ApexAdamW.value):
          moment *= beta1
          moment += (1.0 - beta1) * grad
          velocity *= beta2
          velocity += (1.0 - beta2) * grad * grad
          update = (moment / beta1_correction) / (
              tl.math.sqrt(velocity / beta2_correction) + eps
          )
          update += weight_decay * param
          param -= lr * update
      elif adam_math_mode == tl.constexpr(
          AdamMathType.PyTorchAdam.value
      ):
          grad += weight_decay * param
          moment *= beta1
          moment += (1.0 - beta1) * grad
          velocity *= beta2
          velocity += (1.0 - beta2) * grad * grad
          # PyTorch computes step_size and denominator separately so it can use addcdiv later.
          step_size = -lr / beta1_correction
          beta2_correction_sqrt = tl.math.sqrt(beta2_correction)
          denom = tl.math.sqrt(velocity) / beta2_correction_sqrt + eps
          param += step_size * (moment / denom)
      else:
          raise ValueError(f"Unknown Adam math mode: {adam_math_mode}")
      return param, moment, velocity
  
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 for calling a global function with example usage: _adam_math(param, grad, moment, velocity, beta1, beta2, beta1_correction, beta2_correction, eps, lr, weight_decay, adam_math_mode) and returns: param,  moment,  velocity