code: def test_projected_normal(shape, dim):
      concentration = torch.randn(shape + (dim,)).requires_grad_()
  
      def model():
          with pyro.plate_stack("plates", shape):
              with pyro.plate("particles", 10000):
                  pyro.sample("x", dist.ProjectedNormal(concentration))
  
      value = poutine.trace(model).get_trace().nodes["x"]["value"]
      assert dist.ProjectedNormal.support.check(value).all()
      expected_probe = get_moments(value)
  
      reparam_model = poutine.reparam(
          model, {"x": ProjectedNormalReparam()}
      )
      value = (
          poutine.trace(reparam_model).get_trace().nodes["x"]["value"]
      )
      assert dist.ProjectedNormal.support.check(value).all()
      actual_probe = get_moments(value)
      assert_close(actual_probe, expected_probe, atol=0.05)
  
      for actual_m, expected_m in zip(actual_probe, expected_probe):
          expected_grad = grad(
              expected_m, [concentration], retain_graph=True
          )
          actual_grad = grad(
              actual_m, [concentration], retain_graph=True
          )
          assert_close(actual_grad, expected_grad, atol=0.1)
  
file_path: pyro/tests/infer/reparam/test_projected_normal.py
imports: import pytest
  import torch
  from torch.autograd import grad
  import pyro
  import pyro.distributions
  from pyro import poutine
  from pyro.infer.reparam import ProjectedNormalReparam
  from tests.common import assert_close
  from .util import check_init_reparam
label_desc: Use image node_id 2 for calling a global function with example usage: test_projected_normal(shape, dim) without return types