code: def dirichlet_categorical(data):
      concentration = torch.tensor([1.0, 1.0, 1.0])
      p_latent = pyro.sample("p", dist.Dirichlet(concentration))
      with pyro.plate("N", data.shape[0]):
          pyro.sample("obs", dist.Categorical(p_latent), obs=data)
      return p_latent
  
file_path: pyro/tests/infer/reparam/test_neutra.py
imports: import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro import optim
  from pyro.distributions.transforms import ComposeTransform
  from pyro.infer import MCMC, NUTS, SVI, Trace_ELBO
  from pyro.infer.autoguide import AutoDiagonalNormal, AutoIAFNormal, AutoMultivariateNormal
  from pyro.infer.mcmc.util import initialize_model
  from pyro.infer.reparam import NeuTraReparam
  from tests.common import assert_close
  from .util import check_init_reparam
label_desc: Use image node_id 2 for calling a global function with example usage: dirichlet_categorical(data) and returns: p_latent