code: def _single_tensor_adam(
      params: List[Tensor],
      grads: List[Tensor],
      exp_avgs: List[Tensor],
      exp_avg_sqs: List[Tensor],
      max_exp_avg_sqs: List[Tensor],
      state_steps: List[Tensor],
      grad_scale: Optional[Tensor],
      found_inf: Optional[Tensor],
      *,
      amsgrad: bool,
      has_complex: bool,
      beta1: float,
      beta2: float,
      lr: Union[float, Tensor],
      weight_decay: float,
      eps: float,
      maximize: bool,
      capturable: bool,
      differentiable: bool,
  ):
      assert grad_scale is None and found_inf is None
  
      if torch.jit.is_scripting():
          # this assert is due to JIT being dumb and not realizing that the ops below
          # have overloads to handle both float and Tensor lrs, so we just assert it's
          # a float since most people using JIT are using floats
          assert isinstance(lr, float)
  
      for i, param in enumerate(params):
          grad = grads[i] if not maximize else -grads[i]
          exp_avg = exp_avgs[i]
          exp_avg_sq = exp_avg_sqs[i]
          step_t = state_steps[i]
  
          # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable]
          if not torch._utils.is_compiling() and capturable:
              assert (param.is_cuda and step_t.is_cuda) or (
                  param.is_xla and step_t.is_xla
              ), "If capturable=True, params and state_steps must be CUDA or XLA tensors."
  
          # update step
          step_t += 1
  
          if weight_decay != 0:
              grad = grad.add(param, alpha=weight_decay)
  
          if torch.is_complex(param):
              grad = torch.view_as_real(grad)
              exp_avg = torch.view_as_real(exp_avg)
              exp_avg_sq = torch.view_as_real(exp_avg_sq)
              if amsgrad:
                  max_exp_avg_sqs[i] = torch.view_as_real(
                      max_exp_avg_sqs[i]
                  )
              param = torch.view_as_real(param)
  
          # Decay the first and second moment running average coefficient
          exp_avg.lerp_(grad, 1 - beta1)
          exp_avg_sq.mul_(beta2).addcmul_(
              grad, grad.conj(), value=1 - beta2
          )
  
          if capturable or differentiable:
              step = step_t
  
              bias_correction1 = 1 - beta1**step
              bias_correction2 = 1 - beta2**step
  
              step_size = lr / bias_correction1
              step_size_neg = step_size.neg()
  
              bias_correction2_sqrt = bias_correction2.sqrt()
  
              if amsgrad:
                  # Maintains the maximum of all 2nd moment running avg. till now
                  if differentiable:
                      max_exp_avg_sq = max_exp_avg_sqs[i].clone()
                  else:
                      max_exp_avg_sq = max_exp_avg_sqs[i]
  
                  max_exp_avg_sqs[i].copy_(
                      torch.maximum(max_exp_avg_sq, exp_avg_sq)
                  )
  
                  # Uses the max. for normalizing running avg. of gradient
                  # Folds in (admittedly ugly) 1-elem step_size math here to avoid extra param-set-sized read+write
                  # (can't fold it into addcdiv_ below because addcdiv_ requires value is a Number, not a Tensor)
                  denom = (
                      max_exp_avg_sqs[i].sqrt()
                      / (bias_correction2_sqrt * step_size_neg)
                  ).add_(eps / step_size_neg)
              else:
                  denom = (
                      exp_avg_sq.sqrt()
                      / (bias_correction2_sqrt * step_size_neg)
                  ).add_(eps / step_size_neg)
  
              param.addcdiv_(exp_avg, denom)
          else:
              step = _get_value(step_t)
  
              bias_correction1 = 1 - beta1**step
              bias_correction2 = 1 - beta2**step
  
              step_size = lr / bias_correction1
  
              bias_correction2_sqrt = _dispatch_sqrt(bias_correction2)
  
              if amsgrad:
                  # Maintains the maximum of all 2nd moment running avg. till now
                  torch.maximum(
                      max_exp_avg_sqs[i],
                      exp_avg_sq,
                      out=max_exp_avg_sqs[i],
                  )
  
                  # Use the max. for normalizing running avg. of gradient
                  denom = (
                      max_exp_avg_sqs[i].sqrt() / bias_correction2_sqrt
                  ).add_(eps)
              else:
                  denom = (
                      exp_avg_sq.sqrt() / bias_correction2_sqrt
                  ).add_(eps)
  
              param.addcdiv_(exp_avg, denom, value=-step_size)
  
          # Lastly, switch back to complex view
          if amsgrad and torch.is_complex(params[i]):
              max_exp_avg_sqs[i] = torch.view_as_complex(
                  max_exp_avg_sqs[i]
              )
  
file_path: pytorch/torch/optim/adam.py
imports: from typing import List, Optional, Union, Tuple
  import torch
  from torch import Tensor
  from .optimizer import Optimizer, ParamsT, _use_grad_for_differentiable, _get_value, _stack_if_compiling, _dispatch_sqrt, _default_to_fused_or_foreach, _get_scalar_dtype, _capturable_doc, _differentiable_doc, _foreach_doc, _fused_doc, _maximize_doc, _view_as_real
  from torch.utils._foreach_utils import _get_fused_kernels_supported_devices
label_desc: Use image node_id 2 for calling a global function with example usage: _single_tensor_adam(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, grad_scale, found_inf) without return types