code: def test_nested():
      shape = (5, 6)
  
      @poutine.reparam(
          config={
              "x": HaarReparam(dim=-1),
              "x_haar": HaarReparam(dim=-2),
          }
      )
      def model():
          pyro.sample(
              "x", dist.Normal(torch.zeros(shape), 1).to_event(2)
          )
  
      # Try without initialization, e.g. in AutoGuide._setup_prototype().
      trace = poutine.trace(model).get_trace()
      assert {"x", "x_haar", "x_haar_haar"}.issubset(trace.nodes)
      assert trace.nodes["x"]["is_observed"]
      assert trace.nodes["x_haar"]["is_observed"]
      assert not trace.nodes["x_haar_haar"]["is_observed"]
      assert trace.nodes["x"]["value"].shape == shape
  
      # Try conditioning on x_haar_haar, e.g. in Predictive.
      x = torch.randn(shape)
      x_haar = HaarTransform(dim=-1)(x)
      x_haar_haar = HaarTransform(dim=-2)(x_haar)
      with poutine.condition(data={"x_haar_haar": x_haar_haar}):
          trace = poutine.trace(model).get_trace()
          assert {"x", "x_haar", "x_haar_haar"}.issubset(trace.nodes)
          assert trace.nodes["x"]["is_observed"]
          assert trace.nodes["x_haar"]["is_observed"]
          assert trace.nodes["x_haar_haar"]["is_observed"]
          assert_close(trace.nodes["x"]["value"], x)
          assert_close(trace.nodes["x_haar"]["value"], x_haar)
          assert_close(trace.nodes["x_haar_haar"]["value"], x_haar_haar)
  
      # Try with custom initialization.
      # This is required for autoguides and MCMC.
      with InitMessenger(init_to_value(values={"x": x})):
          trace = poutine.trace(model).get_trace()
          assert {"x", "x_haar", "x_haar_haar"}.issubset(trace.nodes)
          assert trace.nodes["x"]["is_observed"]
          assert trace.nodes["x_haar"]["is_observed"]
          assert not trace.nodes["x_haar_haar"]["is_observed"]
          assert_close(trace.nodes["x"]["value"], x)
  
      # Try conditioning on x.
      x = torch.randn(shape)
      with poutine.condition(data={"x": x}):
          trace = poutine.trace(model).get_trace()
          assert {"x", "x_haar", "x_haar_haar"}.issubset(trace.nodes)
          assert trace.nodes["x"]["is_observed"]
          assert trace.nodes["x_haar"]["is_observed"]
          # TODO Decide whether it is worth fixing this failing assertion.
          # See https://github.com/pyro-ppl/pyro/issues/2878
          # assert trace.nodes["x_haar_haar"]["is_observed"]
          assert_close(trace.nodes["x"]["value"], x)
  
file_path: pyro/tests/infer/reparam/test_haar.py
imports: import pytest
  import torch
  from torch.autograd import grad
  import pyro
  import pyro.distributions
  from pyro import poutine
  from pyro.distributions.transforms.haar import HaarTransform
  from pyro.infer.autoguide.initialization import InitMessenger, init_to_value
  from pyro.infer.reparam import HaarReparam
  from tests.common import assert_close
  from .util import check_init_reparam
label_desc: Use image node_id 5 for calling a global function with example usage: test_nested() without return types