code: def test_dirichlet_categorical_grad_adapt():
      def model(data):
          concentration = torch.tensor([1.0, 1.0, 1.0])
          p_latent = pyro.sample(
              "p_latent", dist.Dirichlet(concentration)
          )
          pyro.sample("obs", dist.Categorical(p_latent), obs=data)
          return p_latent
  
      true_probs = torch.tensor([0.1, 0.6, 0.3])
      data = dist.Categorical(true_probs).sample(
          sample_shape=(torch.Size((2000,)))
      )
      nuts_kernel = NUTS(
          model, jit_compile=True, ignore_jit_warnings=True
      )
      nuts_kernel.mass_matrix_adapter = ArrowheadMassMatrix()
      mcmc = MCMC(nuts_kernel, num_samples=200, warmup_steps=100)
      mcmc.run(data)
      samples = mcmc.get_samples()
      posterior = samples["p_latent"]
      assert_equal(posterior.mean(0), true_probs, prec=0.02)
  
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 16 for calling a global function with example usage: test_dirichlet_categorical_grad_adapt() and returns: p_latent