code: def compute_eigenvalue(self, module, device=None, scale=1.0):
      block_eigenvalue = []
      param_keys = []
      layers = self.get_layers(module)
  
      for block in range(self.layer_num):
          model_block = layers[block]
  
          # We found this randn() has obvious accuracy impact in some cases, save/recover random state here.
          rng_state = torch.random.get_rng_state()
          if device is None:
              v = [
                  torch.randn(p.size())
                  for p in model_block.parameters()
                  if p.grad is not None and p.grad.grad_fn is not None
              ]
          else:
              v = [
                  torch.randn(p.size(), device=device)
                  for p in model_block.parameters()
                  if p.grad is not None and p.grad.grad_fn is not None
              ]
          torch.random.set_rng_state(rng_state)
  
          grads = [
              param.grad
              for param in model_block.parameters()
              if param.grad is not None
              and param.grad.grad_fn is not None
          ]
          params = [
              param
              for param in model_block.parameters()
              if param.grad is not None
              and param.grad.grad_fn is not None
          ]
  
          layer_keys = [id(p) for p in model_block.parameters()]
          param_keys.append(layer_keys)
  
          v = self.normalize(v)
  
          # Disable eigenvalue if the model doesn't support second order gradients computation,
          # e.g. when enabling DS transformer kernel.
          if len(grads) == 0 or len(params) == 0:
              log_dist(
                  f"The model does NOT support eigenvalue computation.",
                  ranks=[0],
                  level=logging.WARNING,
              )
              return []
  
          i = 0
          eigenvalue_current, eigenvalue_previous = 1.0, 0.0
  
          while (
              (i < self.max_iter)
              and abs(eigenvalue_current) > 0
              and (
                  abs(
                      (eigenvalue_current - eigenvalue_previous)
                      / eigenvalue_current
                  )
                  >= self.tol
              )
          ):  # test convergence criteria
              eigenvalue_previous = eigenvalue_current
  
              Hv = torch.autograd.grad(
                  grads,
                  params,
                  grad_outputs=v,
                  only_inputs=True,
                  retain_graph=True,
              )
              # Hv = [hv.float() for hv in Hv]
              Hv = [self.nan_to_num(hv).float() for hv in Hv]
  
              eigenvalue_current = self.inner_product(Hv, v).item()
  
              v = self.normalize(Hv)
              v = [x / scale for x in v]
              i += 1
  
          eigenvalue_current *= scale
          block_eigenvalue.append(eigenvalue_current)
  
          if self.verbose:
              log_dist(
                  f"block: {block}, power iteration: {i}, eigenvalue: {eigenvalue_current}",
                  ranks=[0],
              )
  
      block_eigenvalue = self.post_process(block_eigenvalue)
  
      if self.verbose:
          log_dist(
              f"post processed block_eigenvalue: {block_eigenvalue}",
              ranks=[0],
          )
  
      # {param_id: (eigenvalue, layer_id)}
      ev_dict = {}
      for i, (layer_keys, value) in enumerate(
          zip(param_keys, block_eigenvalue)
      ):
          ev_dict.update(dict.fromkeys(layer_keys, (value, i)))
  
      return ev_dict
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/runtime/eigenvalue.py
imports: import torch
  from deepspeed.utils import log_dist
  import numpy
  import logging
label_desc: Use image node_id 6 for calling the Eigenvalue obj's underlying member method code with example usage: obj.compute_eigenvalue(module, device, scale) and returns: ev_dict, list