code: def _multi_tensor_nadamw(
      params: List[Tensor],
      grads: List[Tensor],
      exp_avgs: List[Tensor],
      exp_avg_sqs: List[Tensor],
      state_steps: List[Tensor],
      *,
      beta1: float,
      beta2: float,
      lr: float,
      weight_decay: float,
      eps: float,
      maximize: bool,
      capturable: bool,
  ):
      if len(params) == 0:
          return
  
      if 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."
  
      if maximize:
          grads = torch._foreach_neg(tuple(grads))  # type: ignore[assignment]
  
      grads = [
          torch.view_as_real(x) if torch.is_complex(x) else x
          for x in grads
      ]
      exp_avgs = [
          torch.view_as_real(x) if torch.is_complex(x) else x
          for x in exp_avgs
      ]
      exp_avg_sqs = [
          torch.view_as_real(x) if torch.is_complex(x) else x
          for x in exp_avg_sqs
      ]
      params = [
          torch.view_as_real(x) if torch.is_complex(x) else x
          for x in params
      ]
  
      # update steps
      torch._foreach_add_(state_steps, 1)
  
      # Perform stepweight decay
      torch._foreach_mul_(params, 1 - lr * weight_decay)
  
      # Decay the first and second moment running average coefficient
      torch._foreach_mul_(exp_avgs, beta1)
      torch._foreach_add_(exp_avgs, grads, alpha=1 - beta1)
  
      torch._foreach_mul_(exp_avg_sqs, beta2)
      torch._foreach_addcmul_(exp_avg_sqs, grads, grads, 1 - beta2)
  
      if capturable:
          # TODO: use foreach_pow if/when foreach_pow is added
          bias_correction1 = [
              torch.pow(beta1, step) for step in state_steps
          ]
          bias_correction2 = [
              torch.pow(beta2, step) for step in 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)
          torch._foreach_neg_(bias_correction1)
          torch._foreach_neg_(bias_correction2)
  
          # foreach_div doesn't allow a scalar as the first arg
          step_size = torch._foreach_div(bias_correction1, lr)
          torch._foreach_reciprocal_(step_size)
          torch._foreach_neg_(step_size)
  
          bias_correction2_sqrt = torch._foreach_sqrt(bias_correction2)
  
          # Only difference between NAdamW and AdamW in this implementation.
          # The official PyTorch implementation of NAdam uses a different algorithm.
          exp_avgs = torch._foreach_mul(exp_avgs, beta1)
          torch._foreach_add_(exp_avgs, grads, alpha=1 - beta1)
  
          exp_avg_sq_sqrt = torch._foreach_sqrt(exp_avg_sqs)
          torch._foreach_div_(
              exp_avg_sq_sqrt,
              torch._foreach_mul(bias_correction2_sqrt, step_size),
          )
          eps_over_step_size = torch._foreach_div(step_size, eps)
          torch._foreach_reciprocal_(eps_over_step_size)
          denom = torch._foreach_add(
              exp_avg_sq_sqrt, eps_over_step_size
          )
  
          torch._foreach_addcdiv_(params, exp_avgs, denom)
      else:
          bias_correction1 = [
              1 - beta1 ** step.item() for step in state_steps
          ]
          bias_correction2 = [
              1 - beta2 ** step.item() for step in state_steps
          ]
  
          step_size = [(lr / bc) * -1 for bc in bias_correction1]
  
          bias_correction2_sqrt = [
              math.sqrt(bc) for bc in bias_correction2
          ]
  
          # Only difference between NAdamW and AdamW in this implementation.
          # The official PyTorch implementation of NAdam uses a different algorithm.
          exp_avgs = torch._foreach_mul(exp_avgs, beta1)
          torch._foreach_add_(exp_avgs, grads, alpha=1 - beta1)
  
          exp_avg_sq_sqrt = torch._foreach_sqrt(exp_avg_sqs)
          torch._foreach_div_(exp_avg_sq_sqrt, bias_correction2_sqrt)
          denom = torch._foreach_add(exp_avg_sq_sqrt, eps)
  
          torch._foreach_addcdiv_(params, exp_avgs, denom, step_size)
  
file_path: pytorch-image-models/timm/optim/nadamw.py
imports: import math
  from typing import List, Optional
  import torch
  from torch import Tensor
label_desc: Use image node_id 3 for calling a global function with example usage: _multi_tensor_nadamw(params, grads, exp_avgs, exp_avg_sqs, state_steps) without return types