code: def test_neals_funnel_smoke(Guide, jit):
      dim = 10
  
      guide = Guide(neals_funnel)
      svi = SVI(
          neals_funnel, guide, optim.Adam({"lr": 1e-10}), Trace_ELBO()
      )
      for _ in range(10):
          svi.step(dim)
  
      neutra = NeuTraReparam(guide.requires_grad_(False))
      model = neutra.reparam(neals_funnel)
      nuts = NUTS(model, jit_compile=jit, ignore_jit_warnings=True)
      mcmc = MCMC(nuts, num_samples=10, warmup_steps=10)
      mcmc.run(dim)
      samples = mcmc.get_samples()
      # XXX: `MCMC.get_samples` adds a leftmost batch dim to all sites, not uniformly at -max_plate_nesting-1;
      # hence the unsqueeze
      transformed_samples = neutra.transform_sample(
          samples["y_shared_latent"].unsqueeze(-2)
      )
      assert "x" in transformed_samples
      assert "y" in transformed_samples
  
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 3 for calling a global function with example usage: test_neals_funnel_smoke(Guide, jit) without return types