class_docstr: A differentiable version of the AdamW optimizer.
  
  This optimizer creates a gradient tape as it updates parameters.
class_docstr_tok: ["A","differentiable","version","of","the","AdamW","optimizer",".","This","optimizer","creates","a","gradient","tape","as","it","updates","parameters","."]
code: def _update(self, grouped_grads: _GroupedGradsType, **kwargs) -> None:
      zipped = zip(self.param_groups, grouped_grads)
      for group_idx, (group, grads) in enumerate(zipped):
          amsgrad = group["amsgrad"]
          beta1, beta2 = group["betas"]
  
          for p_idx, (p, g) in enumerate(zip(group["params"], grads)):
              if g is None:
                  continue
  
              # Perform stepweight decay
              p = p * (1 - group["lr"] * group["weight_decay"])
  
              if g.is_sparse:
                  raise RuntimeError(
                      "AdamW does not support sparse gradients"
                  )
  
              state = self.state[group_idx][p_idx]
  
              # State initialization
              if len(state) == 0:
                  state["step"] = 0
                  # Exponential moving average of gradient values
                  state["exp_avg"] = _torch.zeros_like(p.data)
                  # Exponential moving average of squared gradient values
                  state["exp_avg_sq"] = _torch.zeros_like(p.data)
                  if amsgrad:
                      # Maintains max of all exp. mov. avg. of sq. grad. vals
                      state["max_exp_avg_sq"] = _torch.zeros_like(
                          p.data
                      )
  
              exp_avg, exp_avg_sq = (
                  state["exp_avg"],
                  state["exp_avg_sq"],
              )
              if amsgrad:
                  max_exp_avg_sq = state["max_exp_avg_sq"]
  
              state["step"] += 1
              bias_correction1 = 1 - beta1 ** state["step"]
              bias_correction2 = 1 - beta2 ** state["step"]
  
              # Decay the first and second moment running average coefficient
              state["exp_avg"] = exp_avg = (exp_avg * beta1) + (
                  1 - beta1
              ) * g
              state["exp_avg_sq"] = exp_avg_sq = (
                  exp_avg_sq * beta2
              ) + (1 - beta2) * g * g
  
              # Deal with stability issues
              mask = exp_avg_sq == 0.0
              _maybe_mask(exp_avg_sq, mask)
  
              if amsgrad:
                  # Maintains the max of all 2nd moment running avg. till now
                  state["max_exp_avg_sq"] = max_exp_avg_sq = _torch.max(
                      max_exp_avg_sq, exp_avg_sq
                  )
                  # Use the max. for normalizing running avg. of gradient
                  denom = _add(
                      max_exp_avg_sq.sqrt()
                      / _math.sqrt(bias_correction2),
                      group["eps"],
                  )
              else:
                  denom = _add(
                      exp_avg_sq.sqrt() / _math.sqrt(bias_correction2),
                      group["eps"],
                  )
  
              step_size = group["lr"] / bias_correction1
  
              group["params"][p_idx] = _addcdiv(
                  p, -step_size, exp_avg, denom
              )
  
imports: import abc
  import collections
  import copy
  import math
  import typing
  import warnings
  import torch
  from .None import patch
  from .None import utils
label_desc: Use image node_id 1 for calling the DifferentiableAdamW obj's underlying member method code with example usage: obj._update(grouped_grads) without return types