class_docstr: Implements Adahessian algorithm.
  It has been proposed in `ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning`.
  Arguments:
      params (iterable): iterable of parameters to optimize or dicts defining
          parameter groups
      lr (float, optional): learning rate (default: 0.15)
      betas (Tuple[float, float], optional): coefficients used for computing
          running averages of gradient and its square (default: (0.9, 0.999))
      eps (float, optional): term added to the denominator to improve
          numerical stability (default: 1e-4)
      weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
      hessian_power (float, optional): Hessian power (default: 1). You can also try 0.5. For some tasks we found this to result in better performance.
      single_gpu (Bool, optional): Do you use distributed training or not "torch.nn.parallel.DistributedDataParallel" (default: True)
docstr: Performs a single optimization step.
  Arguments:
      gradsH: The gradient used to compute Hessian vector product.
      closure (callable, optional): A closure that reevaluates the model
          and returns the loss.
docstr_tok: ["Performs","a","single","optimization","step",".","Arguments",":","gradsH",":","The","gradient","used","to","compute","Hessian","vector","product",".","closure","(","callable",",","optional",")",":","A","closure","that","reevaluates","the","model","and","returns","the","loss","."]
code: def step(self, closure=None):
      """Performs a single optimization step.
      Arguments:
          gradsH: The gradient used to compute Hessian vector product.
          closure (callable, optional): A closure that reevaluates the model
              and returns the loss.
      """
      loss = None
      if closure is not None:
          loss = closure()
  
      params = []
      groups = []
      grads = []
  
      # Flatten groups into lists, so that
      #  hut_traces can be called with lists of parameters
      #  and grads
      for group in self.param_groups:
          for p in group["params"]:
              if p.grad is not None:
                  params.append(p)
                  groups.append(group)
                  grads.append(p.grad)
  
      # get the Hessian diagonal
  
      hut_traces = self.get_trace(params, grads)
  
      for p, group, grad, hut_trace in zip(
          params, groups, grads, hut_traces
      ):
          state = self.state[p]
  
          # 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 Hessian diagonal square values
              state["exp_hessian_diag_sq"] = torch.zeros_like(p.data)
  
          exp_avg, exp_hessian_diag_sq = (
              state["exp_avg"],
              state["exp_hessian_diag_sq"],
          )
  
          beta1, beta2 = group["betas"]
  
          state["step"] += 1
  
          # Decay the first and second moment running average coefficient
          exp_avg.mul_(beta1).add_(grad.detach_(), alpha=1 - beta1)
          exp_hessian_diag_sq.mul_(beta2).addcmul_(
              hut_trace, hut_trace, value=1 - beta2
          )
  
          bias_correction1 = 1 - beta1 ** state["step"]
          bias_correction2 = 1 - beta2 ** state["step"]
  
          # make the square root, and the Hessian power
          k = group["hessian_power"]
          denom = (
              (exp_hessian_diag_sq.sqrt() ** k)
              / math.sqrt(bias_correction2) ** k
          ).add_(group["eps"])
  
          # make update
          p.data = p.data - group["lr"] * (
              exp_avg / bias_correction1 / denom
              + group["weight_decay"] * p.data
          )
  
      return loss
  
file_path: bert-ultrafast/training/cramming/backend/optimizers/adahessian.py
imports: import math
  import torch
  from torch.optim.optimizer import Optimizer
label_desc: Use image node_id 3 for calling the Adahessian obj's underlying member method code with example usage: obj.step(closure) and returns: loss