code: def test_auto_dirichlet(auto_class, Elbo):
      num_steps = 2000
      prior = torch.tensor([0.5, 1.0, 1.5, 3.0])
      data = torch.tensor([0] * 4 + [1] * 2 + [2] * 5).long()
      posterior = torch.tensor([4.5, 3.0, 6.5, 3.0])
  
      def model(data):
          p = pyro.sample("p", dist.Dirichlet(prior))
          with pyro.plate("data_plate"):
              pyro.sample(
                  "data",
                  dist.Categorical(p).expand_by(data.shape),
                  obs=data,
              )
  
      guide = auto_class(model)
      svi = SVI(model, guide, optim.Adam({"lr": 0.003}), loss=Elbo())
  
      for _ in range(num_steps):
          loss = svi.step(data)
          assert np.isfinite(loss), loss
  
      expected_mean = posterior / posterior.sum()
      if isinstance(guide, (AutoIAFNormal, AutoNormalizingFlow)):
          loc = guide.transform(torch.zeros(guide.latent_dim))
      else:
          loc = guide.loc
      actual_mean = biject_to(constraints.simplex)(loc)
      assert_equal(
          actual_mean,
          expected_mean,
          prec=0.2,
          msg="".join(
              [
                  "\nexpected {}".format(
                      expected_mean.detach().cpu().numpy()
                  ),
                  "\n  actual {}".format(
                      actual_mean.detach().cpu().numpy()
                  ),
              ]
          ),
      )
  
imports: import logging
  from functools import partial
  import numpy
  import pytest
  import torch
  from torch.distributions import biject_to, constraints
  import pyro
  import pyro.distributions
  import pyro.optim
  from pyro.distributions.transforms import block_autoregressive, iterated
  from pyro.infer import SVI, Trace_ELBO, TraceMeanField_ELBO
  from pyro.infer.autoguide import AutoDiagonalNormal, AutoIAFNormal, AutoLaplaceApproximation, AutoLowRankMultivariateNormal, AutoMultivariateNormal
  from pyro.infer.autoguide.guides import AutoNormalizingFlow
  from tests.common import assert_equal
  from tests.integration_tests.test_conjugate_gaussian_models import GaussianChain
label_desc: Use image node_id 3 for calling a global function with example usage: test_auto_dirichlet(auto_class, Elbo) without return types