fix loc scale reparam with center=1 (#1059)

pyro-ppl · Jun 8, 2021 · f212f32 · f212f32
1 parent 4ad6b98
commit f212f32
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/numpyro/infer/reparam.py b/numpyro/infer/reparam.py
@@ -89,7 +89,7 @@ def __call__(self, name, fn, obs):
         assert obs is None, "LocScaleReparam does not support observe statements"
         centered = self.centered
         if is_identically_one(centered):
-            return name, fn, obs
+            return fn, obs
         event_shape = fn.event_shape
         fn, expand_shape, event_dim = self._unwrap(fn)
 

diff --git a/test/infer/test_reparam.py b/test/infer/test_reparam.py
@@ -217,9 +217,9 @@ def get_expected_probe(loc, scale):
         return get_moments(trace["x"]["value"])
 
     if "dist_type" == "Normal":
-        reparam = LocScaleReparam()
+        reparam = LocScaleReparam(centered)
     else:
-        reparam = LocScaleReparam(shape_params=["df"])
+        reparam = LocScaleReparam(centered, shape_params=["df"])
 
     def get_actual_probe(loc, scale):
         with numpyro.handlers.trace() as trace: