code: def _multi_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,
  ):
      if len(params) == 0:
          return
  
      if isinstance(lr, Tensor) and not capturable:
          raise RuntimeError(
              "lr as a Tensor is not supported for capturable=False and foreach=True"
          )
  
      # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable]
      if not torch._utils.is_compiling() and capturable:
          assert all(
              p.is_cuda and step.is_cuda
              for p, step in zip(params, state_steps)
          ), "If capturable=True, params and state_steps must be CUDA tensors."
  
      assert grad_scale is None and found_inf is None
  
      assert not differentiable, "_foreach ops don't support autograd"
  
      grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
          [
              params,
              grads,
              exp_avgs,
              exp_avg_sqs,
              max_exp_avg_sqs,
              state_steps,
          ]
      )
      for (
          device_params,
          device_grads,
          device_exp_avgs,
          device_exp_avg_sqs,
          device_max_exp_avg_sqs,
          device_state_steps,
      ), _ in grouped_tensors.values():
          if maximize:
              device_grads = torch._foreach_neg(device_grads)
  
          # Handle complex parameters
          if has_complex:
              if amsgrad:
                  _view_as_real(
                      device_params,
                      device_grads,
                      device_exp_avgs,
                      device_exp_avg_sqs,
                      device_max_exp_avg_sqs,
                  )
              else:
                  _view_as_real(
                      device_params,
                      device_grads,
                      device_exp_avgs,
                      device_exp_avg_sqs,
                  )
  
          # Update steps
          # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over
          # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just
          # wrapped it once now. The alpha is required to assure we go to the right overload.
          if device_state_steps[0].is_cpu:
              torch._foreach_add_(
                  device_state_steps,
                  torch.tensor(1.0, device="cpu"),
                  alpha=1.0,
              )
          else:
              torch._foreach_add_(device_state_steps, 1)
  
          if weight_decay != 0:
              # Re-use the intermediate memory (device_grads) already allocated for maximize
              if maximize:
                  torch._foreach_add_(
                      device_grads, device_params, alpha=weight_decay
                  )
              else:
                  device_grads = torch._foreach_add(
                      device_grads, device_params, alpha=weight_decay
                  )
  
          # Decay the first and second moment running average coefficient
          torch._foreach_lerp_(device_exp_avgs, device_grads, 1 - beta1)
  
          torch._foreach_mul_(device_exp_avg_sqs, beta2)
          torch._foreach_addcmul_(
              device_exp_avg_sqs, device_grads, device_grads, 1 - beta2
          )
  
          # Delete the local intermediate since it won't be used anymore to save on peak memory
          del device_grads
  
          if capturable:
              bias_correction1 = torch._foreach_pow(
                  beta1, device_state_steps
              )
              bias_correction2 = torch._foreach_pow(
                  beta2, device_state_steps
              )
              # foreach_sub doesn't allow a scalar as the first arg
              torch._foreach_sub_(bias_correction1, 1)
              torch._foreach_sub_(bias_correction2, 1)
              # we do not negate bias_correction1 as it'll need to be negated later anyway
              torch._foreach_neg_(bias_correction2)
  
              # foreach_div doesn't allow a scalar as the first arg
              torch._foreach_div_(bias_correction1, lr)
              torch._foreach_reciprocal_(bias_correction1)
  
              torch._foreach_sqrt_(bias_correction2)
  
              # Re-assign for clarity as we maintain minimal intermediates: we'll have
              # step_size = - lr / (1 - beta1 ^ t) where t = num_steps
              # bias_correction2_sqrt = sqrt(1 - beta2 ^ t)
              step_size = bias_correction1
              bias_correction2_sqrt = bias_correction2
  
              if amsgrad:
                  # Maintains the maximum of all 2nd moment running avg. till now
                  torch._foreach_maximum_(device_max_exp_avg_sqs, device_exp_avg_sqs)  # type: ignore[assignment]
  
                  # Set intermediate to the max. for normalizing running avg. of gradient when amsgrad
                  exp_avg_sq_sqrt = torch._foreach_sqrt(
                      device_max_exp_avg_sqs
                  )
              else:
                  exp_avg_sq_sqrt = torch._foreach_sqrt(
                      device_exp_avg_sqs
                  )
  
              torch._foreach_div_(
                  exp_avg_sq_sqrt, bias_correction2_sqrt
              )
              torch._foreach_add_(exp_avg_sq_sqrt, eps)
              torch._foreach_div_(exp_avg_sq_sqrt, step_size)
  
              # at this point, exp_avg_sq_sqrt = - (1 - beta^t) * [sqrt(exp_avg_sq / (1 - beta2^t)) + eps] / lr
              torch._foreach_addcdiv_(
                  device_params, device_exp_avgs, exp_avg_sq_sqrt
              )
          else:
              bias_correction1 = [
                  1 - beta1 ** _get_value(step)
                  for step in device_state_steps
              ]
              bias_correction2 = [
                  1 - beta2 ** _get_value(step)
                  for step in device_state_steps
              ]
  
              step_size = _stack_if_compiling(
                  [(lr / bc) * -1 for bc in bias_correction1]
              )
  
              bias_correction2_sqrt = [
                  _dispatch_sqrt(bc) for bc in bias_correction2
              ]
  
              if amsgrad:
                  # Maintains the maximum of all 2nd moment running avg. till now
                  torch._foreach_maximum_(
                      device_max_exp_avg_sqs, device_exp_avg_sqs
                  )
  
                  # Use the max. for normalizing running avg. of gradient
                  exp_avg_sq_sqrt = torch._foreach_sqrt(
                      device_max_exp_avg_sqs
                  )
              else:
                  exp_avg_sq_sqrt = torch._foreach_sqrt(
                      device_exp_avg_sqs
                  )
  
              torch._foreach_div_(
                  exp_avg_sq_sqrt, bias_correction2_sqrt
              )
              torch._foreach_add_(exp_avg_sq_sqrt, eps)
              torch._foreach_addcdiv_(
                  device_params,
                  device_exp_avgs,
                  exp_avg_sq_sqrt,
                  step_size,
              )
  
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 3 for calling a global function with example usage: _multi_tensor_adam(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, grad_scale, found_inf) without return types