code: def test_arrowhead_mass_matrix(diagonal):
      shape = (2, 3)
      num_samples = 1000
  
      size = shape[0] * shape[1]
      block_adapter = BlockMassMatrix()
      arrowhead_adapter = ArrowheadMassMatrix()
      mass_matrix_shape = (size,) if diagonal else (size, size)
      block_adapter.configure({("z",): mass_matrix_shape})
      arrowhead_adapter.configure({("z",): mass_matrix_shape})
  
      cov = torch.randn(size, size)
      cov = torch.mm(cov, cov.t())
      if diagonal:
          cov = cov.diag().diag()
      z_dist = torch.distributions.MultivariateNormal(
          torch.zeros(size), covariance_matrix=cov
      )
      g_dist = torch.distributions.MultivariateNormal(
          torch.zeros(size), precision_matrix=cov
      )
      z_samples = z_dist.sample((num_samples,)).reshape(
          (num_samples,) + shape
      )
      g_samples = g_dist.sample((num_samples,)).reshape(
          (num_samples,) + shape
      )
  
      for i in range(num_samples):
          block_adapter.update({"z": z_samples[i]}, {"z": g_samples[i]})
          arrowhead_adapter.update(
              {"z": z_samples[i]}, {"z": g_samples[i]}
          )
      block_adapter.end_adaptation()
      arrowhead_adapter.end_adaptation()
  
      assert_close(
          arrowhead_adapter.inverse_mass_matrix[("z",)],
          block_adapter.inverse_mass_matrix[("z",)],
          atol=0.3,
          rtol=0.3,
      )
  
imports: import pytest
  import torch
  from pyro.infer.mcmc.adaptation import ArrowheadMassMatrix, BlockMassMatrix, WarmupAdapter, adapt_window
  from tests.common import assert_close, assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: test_arrowhead_mass_matrix(diagonal) without return types