class_docstr: EXPERIMENTAL This class is used to adapt (inverse) mass matrix and provide useful
  methods to calculate algebraic terms which involves the mass matrix.
  
  The mass matrix will have arrowhead structure, with the head including all
  dense sites specified in the argument `full_mass` of the HMC/NUTS kernels.
  
  :param float init_scale: initial scale to construct the initial mass matrix.
class_docstr_tok: ["EXPERIMENTAL","This","class","is","used","to","adapt","(","inverse",")","mass","matrix","and","provide","useful","methods","to","calculate","algebraic","terms","which","involves","the","mass","matrix",".","The","mass","matrix","will","have","arrowhead","structure",",","with","the","head","including","all","dense","sites","specified","in","the","argument","`","full_mass","`","of","the","HMC\/NUTS","kernels",".",":","param","float","init_scale",":","initial","scale","to","construct","the","initial","mass","matrix","."]
docstr: Computes the gradient of kinetic energy w.r.t. the momentum `r`.
  It is equivalent to compute velocity given the momentum `r`.
  
  :param dict r: a dictionary maps site names to a tensor momentum.
  :returns: a dictionary maps site names to the corresponding gradient
docstr_tok: ["Computes","the","gradient","of","kinetic","energy","w.r.t",".","the","momentum","`","r","`",".","It","is","equivalent","to","compute","velocity","given","the","momentum","`","r","`",".",":","param","dict","r",":","a","dictionary","maps","site","names","to","a","tensor","momentum",".",":","returns",":","a","dictionary","maps","site","names","to","the","corresponding","gradient"]
code: def kinetic_grad(self, r):
      """
      Computes the gradient of kinetic energy w.r.t. the momentum `r`.
      It is equivalent to compute velocity given the momentum `r`.
  
      :param dict r: a dictionary maps site names to a tensor momentum.
      :returns: a dictionary maps site names to the corresponding gradient
      """
      v = {}
      for (
          site_names,
          mass_matrix_sqrt_inverse,
      ) in self._mass_matrix_sqrt_inverse.items():
          r_flat = torch.cat(
              [r[site_name].reshape(-1) for site_name in site_names]
          )
          # NB: using inverse_mass_matrix as in BlockMassMatrix will cost
          # O(N^2 x head_size) operators and O(N^2) memory requirement;
          # here, we will leverage mass_matrix_sqrt_inverse to reduce the cost to
          # O(N x head_size^2) operators and O(N x head_size) memory requirement.
          r_unscaled = triu_matvecmul(mass_matrix_sqrt_inverse, r_flat)
          v_flat = triu_matvecmul(
              mass_matrix_sqrt_inverse, r_unscaled, transpose=True
          )
  
          # unpacking
          pos = 0
          for site_name in site_names:
              next_pos = pos + r[site_name].numel()
              v[site_name] = v_flat[pos:next_pos].reshape(
                  r[site_name].shape
              )
              pos = next_pos
      return v
  
imports: import math
  from collections import namedtuple
  import torch
  import pyro
  from pyro.ops.arrowhead import SymmArrowhead, sqrt, triu_gram, triu_inverse, triu_matvecmul
  from pyro.ops.dual_averaging import DualAveraging
  from pyro.ops.welford import WelfordArrowheadCovariance, WelfordCovariance
label_desc: Use image node_id 9 for calling the ArrowheadMassMatrix obj's underlying member method code with example usage: obj.kinetic_grad(r) and returns: v