code: def test_structured_mass():
      def model(cov):
          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([1]).to_event(1)
          )
          wxyz = torch.cat([w, x, y, z])
          pyro.sample(
              "obs",
              dist.MultivariateNormal(torch.zeros(5), cov),
              obs=wxyz,
          )
  
      w_cov = torch.tensor([[1.5, 0.5], [0.5, 1.5]])
      xy_cov = torch.tensor([[2.0, 1.0], [1.0, 3.0]])
      z_var = torch.tensor([2.5])
      cov = torch.zeros(5, 5)
      cov[:2, :2] = w_cov
      cov[2:4, 2:4] = xy_cov
      cov[4, 4] = z_var
  
      # 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(cov)
          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",), ("x", "y")],
      )
      mcmc = MCMC(kernel, num_samples=1, warmup_steps=1000)
      mcmc.run(cov)
      assert_close(
          kernel.inverse_mass_matrix[("w",)], w_cov, atol=0.5, rtol=0.5
      )
      assert_close(
          kernel.inverse_mass_matrix[("x", "y")],
          xy_cov,
          atol=0.5,
          rtol=0.5,
      )
      assert_close(
          kernel.inverse_mass_matrix[("z",)], z_var, atol=0.5, rtol=0.5
      )
  
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 14 for calling a global function with example usage: test_structured_mass() without return types