code: def test_tmc_categoricals(
      depth, max_plate_nesting, num_samples, tmc_strategy
  ):
      qs = [
          pyro.param("q0", torch.tensor([0.4, 0.6], requires_grad=True))
      ]
      for i in range(1, depth):
          qs.append(
              pyro.param(
                  "q{}".format(i),
                  torch.randn(2, 2).abs().detach().requires_grad_(),
                  constraint=constraints.simplex,
              )
          )
      qs.append(
          pyro.param(
              "qy", torch.tensor([0.75, 0.25], requires_grad=True)
          )
      )
  
      qs = [q.unconstrained() for q in qs]
  
      data = (torch.rand(4, 3) > 0.5).to(
          dtype=qs[-1].dtype, device=qs[-1].device
      )
  
      def model():
          x = pyro.sample("x0", dist.Categorical(pyro.param("q0")))
          with pyro.plate("local", 3):
              for i in range(1, depth):
                  x = pyro.sample(
                      "x{}".format(i),
                      dist.Categorical(
                          pyro.param("q{}".format(i))[..., x, :]
                      ),
                  )
              with pyro.plate("data", 4):
                  pyro.sample(
                      "y",
                      dist.Bernoulli(pyro.param("qy")[..., x]),
                      obs=data,
                  )
  
      elbo = TraceEnum_ELBO(max_plate_nesting=max_plate_nesting)
      enum_model = config_enumerate(
          model,
          default="parallel",
          expand=False,
          num_samples=None,
          tmc=tmc_strategy,
      )
      expected_loss = (
          -elbo.differentiable_loss(enum_model, lambda: None)
      ).exp()
      expected_grads = grad(expected_loss, qs)
  
      tmc = TraceTMC_ELBO(max_plate_nesting=max_plate_nesting)
      tmc_model = config_enumerate(
          model,
          default="parallel",
          expand=False,
          num_samples=num_samples,
          tmc=tmc_strategy,
      )
      actual_loss = (
          -tmc.differentiable_loss(tmc_model, lambda: None)
      ).exp()
      actual_grads = grad(actual_loss, qs)
  
      prec = 0.05
      assert_equal(
          actual_loss,
          expected_loss,
          prec=prec,
          msg="".join(
              [
                  "\nexpected loss = {}".format(expected_loss),
                  "\n  actual loss = {}".format(actual_loss),
              ]
          ),
      )
  
      for actual_grad, expected_grad in zip(
          actual_grads, expected_grads
      ):
          assert_equal(
              actual_grad,
              expected_grad,
              prec=prec,
              msg="".join(
                  [
                      "\nexpected grad = {}".format(
                          expected_grad.detach().cpu().numpy()
                      ),
                      "\n  actual grad = {}".format(
                          actual_grad.detach().cpu().numpy()
                      ),
                  ]
              ),
          )
  
file_path: pyro/tests/infer/test_tmc.py
imports: import logging
  import math
  import pytest
  import torch
  from torch.autograd import grad
  from torch.distributions import constraints
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.distributions.testing import fakes
  from pyro.infer import config_enumerate
  from pyro.infer.importance import vectorized_importance_weights
  from pyro.infer.traceenum_elbo import TraceEnum_ELBO
  from pyro.infer.tracetmc_elbo import TraceTMC_ELBO
  from tests.common import assert_equal
label_desc: Use image node_id 1 for calling a global function with example usage: test_tmc_categoricals(depth, max_plate_nesting, num_samples, tmc_strategy) without return types