code: def test_moments(dist_type, centered, shape):
      loc = torch.empty(shape).uniform_(-1.0, 1.0).requires_grad_()
      scale = torch.empty(shape).uniform_(0.5, 1.5).requires_grad_()
      if isinstance(centered, torch.Tensor):
          centered = centered.expand(shape)
  
      def model():
          with pyro.plate_stack("plates", shape):
              with pyro.plate("particles", 200000):
                  if "dist_type" == "Normal":
                      pyro.sample("x", dist.Normal(loc, scale))
                  elif "dist_type" == "StudentT":
                      pyro.sample("x", dist.StudentT(10.0, loc, scale))
                  else:
                      pyro.sample(
                          "x", dist.AsymmetricLaplace(loc, scale, 1.5)
                      )
  
      value = poutine.trace(model).get_trace().nodes["x"]["value"]
      expected_probe = get_moments(value)
  
      reparam = LocScaleReparam(centered)
      reparam_model = poutine.reparam(model, {"x": reparam})
      value = (
          poutine.trace(reparam_model).get_trace().nodes["x"]["value"]
      )
      actual_probe = get_moments(value)
  
      if not is_identically_one(centered):
          if "dist_type" == "Normal":
              assert reparam.shape_params == ()
          elif "dist_type" == "StudentT":
              assert reparam.shape_params == ("df",)
          else:
              assert reparam.shape_params == ("asymmetry",)
  
      assert_close(actual_probe, expected_probe, atol=0.1, rtol=0.05)
  
      for actual_m, expected_m in zip(actual_probe, expected_probe):
          expected_grads = grad(
              expected_m.sum(), [loc, scale], retain_graph=True
          )
          actual_grads = grad(
              actual_m.sum(), [loc, scale], retain_graph=True
          )
          assert_close(
              actual_grads[0], expected_grads[0], atol=0.1, rtol=0.05
          )
          assert_close(
              actual_grads[1], expected_grads[1], atol=0.1, rtol=0.05
          )
  
file_path: pyro/tests/infer/reparam/test_loc_scale.py
imports: import pytest
  import torch
  from torch.autograd import grad
  import pyro
  import pyro.distributions
  from pyro import poutine
  from pyro.distributions.util import is_identically_one
  from pyro.infer.autoguide.initialization import InitMessenger, init_to_uniform
  from pyro.infer.reparam import LocScaleReparam
  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_moments(dist_type, centered, shape) without return types