code: def test_arrowhead_mass():
      def model(prec):
          w = pyro.sample(
              "w", dist.Normal(0, 1000).expand([2]).to_event(1)
          )
          x = pyro.sample(
              "x", dist.Normal(0, 1000).expand([1]).to_event(1)
          )
          y = pyro.sample(
              "y", dist.Normal(0, 1000).expand([1]).to_event(1)
          )
          z = pyro.sample(
              "z", dist.Normal(0, 1000).expand([2]).to_event(1)
          )
          wyxz = torch.cat([w, y, x, z])
          pyro.sample(
              "obs",
              dist.MultivariateNormal(
                  torch.zeros(6), precision_matrix=prec
              ),
              obs=wyxz,
          )
  
      A = torch.randn(6, 12)
      prec = A @ A.t() * 0.1
  
      # smoke tests
      for dense_mass in [True, False]:
          kernel = NUTS(
              model,
              jit_compile=True,
              ignore_jit_warnings=True,
              full_mass=dense_mass,
          )
          mcmc = MCMC(kernel, num_samples=1, warmup_steps=1)
          mcmc.run(prec)
          assert kernel.inverse_mass_matrix[
              ("w", "x", "y", "z")
          ].dim() == 1 + int(dense_mass)
  
      kernel = NUTS(
          model,
          jit_compile=True,
          ignore_jit_warnings=True,
          full_mass=[("w",), ("y", "x")],
      )
      kernel.mass_matrix_adapter = ArrowheadMassMatrix()
      mcmc = MCMC(kernel, num_samples=1, warmup_steps=1000)
      mcmc.run(prec)
      assert ("w", "y", "x", "z") in kernel.inverse_mass_matrix
      mass_matrix = kernel.mass_matrix_adapter.mass_matrix[
          ("w", "y", "x", "z")
      ]
      assert mass_matrix.top.shape == (4, 6)
      assert mass_matrix.bottom_diag.shape == (2,)
      assert_close(mass_matrix.top, prec[:4], atol=0.2, rtol=0.2)
      assert_close(
          mass_matrix.bottom_diag, prec.diag()[4:], atol=0.2, rtol=0.2
      )
  
imports: import logging
  import os
  from collections import namedtuple
  import pytest
  import torch
  import pyro
  import pyro.distributions
  import pyro.optim
  import pyro.poutine
  from pyro.contrib.conjugate.infer import BetaBinomialPair, GammaPoissonPair, collapse_conjugate, posterior_replay
  from pyro.infer import SVI, TraceEnum_ELBO
  from pyro.infer.autoguide import AutoDelta
  from pyro.infer.mcmc import MCMC, NUTS, ArrowheadMassMatrix
  from pyro.util import ignore_jit_warnings
  from tests.common import assert_close, assert_equal
  from .test_hmc import GaussianChain, rmse
label_desc: Use image node_id 15 for calling a global function with example usage: test_arrowhead_mass() without return types