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: 0.5)
      seed (int, optional): Random number generator seed (default: None)
  
      Example:
      >>> import torch_optimizer as optim
      >>> optimizer = optim.Adahessian(model.parameters(), lr = 1.0)
      >>> optimizer.zero_grad()
      >>> loss_fn(model(input), target).backward(create_graph=True)
      >>> optimizer.step()
  
      __ https://arxiv.org/abs/2006.00719
  
      Note:
          Reference code: https://github.com/amirgholami/adahessian
docstr: Get an estimate of Hessian Trace.
  This is done by computing the Hessian vector product with a random
  vector v at the current gradient point, to estimate Hessian trace by
  computing the gradient of <gradsH,v>.
  :param gradsH: a list of torch variables
  :return: a list of torch tensors
docstr_tok: ["Get","an","estimate","of","Hessian","Trace",".","This","is","done","by","computing","the","Hessian","vector","product","with","a","random","vector","v","at","the","current","gradient","point",",","to","estimate","Hessian","trace","by","computing","the","gradient","of","<","gradsH",",","v",">",".",":","param","gradsH",":","a","list","of","torch","variables",":","return",":","a","list","of","torch","tensors"]
code: def get_trace(
      self, params: Params, grads: Grads
  ) -> List[torch.Tensor]:
      """Get an estimate of Hessian Trace.
      This is done by computing the Hessian vector product with a random
      vector v at the current gradient point, to estimate Hessian trace by
      computing the gradient of <gradsH,v>.
      :param gradsH: a list of torch variables
      :return: a list of torch tensors
      """
  
      # Check backward was called with create_graph set to True
      for i, grad in enumerate(grads):
          if grad.grad_fn is None:
              msg = (
                  "Gradient tensor {:} does not have grad_fn. When "
                  "calling loss.backward(), make sure the option "
                  "create_graph is set to True."
              )
              raise RuntimeError(msg.format(i))
  
      v = [
          2
          * torch.randint_like(
              p, high=2, memory_format=torch.preserve_format
          )
          - 1
          for p in params
      ]
  
      # this is for distributed setting with single node and multi-gpus,
      # for multi nodes setting, we have not support it yet.
      hvs = torch.autograd.grad(
          grads,
          params,
          grad_outputs=v,
          only_inputs=True,
          retain_graph=True,
      )
  
      hutchinson_trace = []
      for hv in hvs:
          param_size = hv.size()
          if len(param_size) <= 2:  # for 0/1/2D tensor
              # Hessian diagonal block size is 1 here.
              # We use that torch.abs(hv * vi) = hv.abs()
              tmp_output = hv.abs()
  
          elif len(param_size) == 4:  # Conv kernel
              # Hessian diagonal block size is 9 here: torch.sum() reduces
              # the dim 2/3.
              # We use that torch.abs(hv * vi) = hv.abs()
              tmp_output = torch.mean(
                  hv.abs(), dim=[2, 3], keepdim=True
              )
          hutchinson_trace.append(tmp_output)
  
      return hutchinson_trace
  
file_path: pytorch-optimizer/torch_optimizer/adahessian.py
imports: import math
  from typing import List, Optional
  import torch
  from torch.optim.optimizer import Optimizer
  from .types import Betas2, OptFloat, OptLossClosure, Params
label_desc: Use image node_id 2 for calling the Adahessian obj's underlying member method code with example usage: obj.get_trace(params, grads) and returns: hutchinson_trace