code: def test_tmc_categoricals(
      depth, max_plate_nesting, num_samples, tmc_strategy
  ):
      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,
                  )
  
      with pyro_backend("pyro"):
          # initialize
          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
          )
  
      with pyro_backend("pyro"):
          elbo = infer.TraceTMC_ELBO(
              max_plate_nesting=max_plate_nesting
          )
          enum_model = infer.config_enumerate(
              model,
              default="parallel",
              expand=False,
              num_samples=num_samples,
              tmc=tmc_strategy,
          )
          expected_loss = (
              -elbo.differentiable_loss(enum_model, lambda: None)
          ).exp()
          expected_grads = grad(expected_loss, qs)
  
      with pyro_backend("contrib.funsor"):
          tmc = infer.TraceTMC_ELBO(max_plate_nesting=max_plate_nesting)
          tmc_model = infer.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/contrib/funsor/test_tmc.py
imports: import logging
  import math
  import pytest
  import torch
  from torch.autograd import grad
  from torch.distributions import constraints
  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