code: def test_tmc_normals_chain_gradient(
      depth,
      num_samples,
      max_plate_nesting,
      expand,
      guide_type,
      reparameterized,
      tmc_strategy,
  ):
      def model(reparameterized):
          Normal = (
              dist.Normal
              if reparameterized
              else dist.testing.fakes.NonreparameterizedNormal
          )
          x = pyro.sample(
              "x0", Normal(pyro.param("q2"), math.sqrt(1.0 / depth))
          )
          for i in range(1, depth):
              x = pyro.sample(
                  "x{}".format(i), Normal(x, math.sqrt(1.0 / depth))
              )
          pyro.sample("y", Normal(x, 1.0), obs=torch.tensor(float(1)))
  
      def factorized_guide(reparameterized):
          Normal = (
              dist.Normal
              if reparameterized
              else dist.testing.fakes.NonreparameterizedNormal
          )
          pyro.sample(
              "x0", Normal(pyro.param("q2"), math.sqrt(1.0 / depth))
          )
          for i in range(1, depth):
              pyro.sample(
                  "x{}".format(i),
                  Normal(0.0, math.sqrt(float(i + 1) / depth)),
              )
  
      def nonfactorized_guide(reparameterized):
          Normal = (
              dist.Normal
              if reparameterized
              else dist.testing.fakes.NonreparameterizedNormal
          )
          x = pyro.sample(
              "x0", Normal(pyro.param("q2"), math.sqrt(1.0 / depth))
          )
          for i in range(1, depth):
              x = pyro.sample(
                  "x{}".format(i), Normal(x, math.sqrt(1.0 / depth))
              )
  
      with pyro_backend("contrib.funsor"):
          # compare reparameterized and nonreparameterized gradient estimates
          q2 = pyro.param("q2", torch.tensor(0.5, requires_grad=True))
          qs = (q2.unconstrained(),)
  
          tmc = infer.TraceTMC_ELBO(max_plate_nesting=max_plate_nesting)
          tmc_model = infer.config_enumerate(
              model,
              default="parallel",
              expand=expand,
              num_samples=num_samples,
              tmc=tmc_strategy,
          )
          guide = (
              factorized_guide
              if guide_type == "factorized"
              else nonfactorized_guide
              if guide_type == "nonfactorized"
              else lambda *args: None
          )
          tmc_guide = infer.config_enumerate(
              guide,
              default="parallel",
              expand=expand,
              num_samples=num_samples,
              tmc=tmc_strategy,
          )
  
          # convert to linear space for unbiasedness
          actual_loss = (
              -tmc.differentiable_loss(
                  tmc_model, tmc_guide, reparameterized
              )
          ).exp()
          actual_grads = grad(actual_loss, qs)
  
      # gold values from Funsor
      expected_grads = (
          torch.tensor(
              {1: 0.0999, 2: 0.0860, 3: 0.0802, 4: 0.0771}[depth]
          ),
      )
  
      grad_prec = 0.05 if reparameterized else 0.1
  
      for actual_grad, expected_grad in zip(
          actual_grads, expected_grads
      ):
          print(actual_loss)
          assert_equal(
              actual_grad,
              expected_grad,
              prec=grad_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 2 for calling a global function with example usage: test_tmc_normals_chain_gradient(depth, num_samples, max_plate_nesting, expand, guide_type, reparameterized, tmc_strategy) without return types