Enhance LocScaleReparam noncenter dependency (#1384)

* Enhance LocScaleReparam noncenter dependency

* fix lint

* Use id in the first place

* change tree_multimap to tree_map

* Fix lint

* Fix lint

* Fix failing tests due to FutureWarning tree_multimap

* Also support ndarray in LocScaleReparam

Author: fehiepsi

numpyro/infer/reparam.py

@@ -4,13 +4,17 @@
 from abc import ABC, abstractmethod
 import math
 
+import numpy as np
+
+import jax
 import jax.numpy as jnp
 
 import numpyro
 import numpyro.distributions as dist
 from numpyro.distributions import biject_to, constraints
 from numpyro.distributions.util import is_identically_one, safe_normalize, sum_rightmost
 from numpyro.infer.autoguide import AutoContinuous
+from numpyro.util import not_jax_tracer
 
 
 class Reparam(ABC):
@@ -78,11 +82,19 @@ class LocScaleReparam(Reparam):
     """
 
     def __init__(self, centered=None, shape_params=()):
-        assert centered is None or isinstance(centered, (int, float))
+        assert centered is None or isinstance(
+            centered, (int, float, np.generic, np.ndarray, jnp.ndarray, jax.core.Tracer)
+        )
         assert isinstance(shape_params, (tuple, list))
         assert all(isinstance(name, str) for name in shape_params)
-        if isinstance(centered, (int, float)):
-            assert 0 <= centered and centered <= 1
+        if centered is not None:
+            is_valid = constraints.unit_interval.check(centered)
+            if not_jax_tracer(is_valid):
+                if not np.all(is_valid):
+                    raise ValueError(
+                        "`centered` argument does not satisfy `0 <= centered <= 1`."
+                    )
+
         self.centered = centered
         self.shape_params = shape_params
 
@@ -102,8 +114,12 @@ def __call__(self, name, fn, obs):
                 jnp.full(event_shape, 0.5),
                 constraint=constraints.unit_interval,
             )
-        params["loc"] = fn.loc * centered
-        params["scale"] = fn.scale**centered
+        if isinstance(centered, (int, float, np.generic)) and centered == 0.0:
+            params["loc"] = jnp.zeros_like(fn.loc)
+            params["scale"] = jnp.ones_like(fn.scale)
+        else:
+            params["loc"] = fn.loc * centered
+            params["scale"] = fn.scale**centered
         decentered_fn = self._wrap(type(fn)(**params), expand_shape, event_dim)
 
         # Draw decentered noise.

test/infer/test_reparam.py

@@ -233,6 +233,8 @@ def test_loc_scale(dist_type, centered, shape, event_dim):
     loc = np.random.uniform(-1.0, 1.0, shape)
     scale = np.random.uniform(0.5, 1.5, shape)
     event_dim = min(event_dim, len(shape))
+    if event_dim == 1 and centered is not None:
+        centered = jnp.broadcast_to(centered, shape[-event_dim:])
 
     def model(loc, scale):
         with numpyro.plate_stack("plates", shape[: len(shape) - event_dim]):
@@ -272,6 +274,22 @@ def get_actual_probe(loc, scale):
     assert_allclose(actual_grad[1], expected_grad[1], atol=0.05)  # scale grad
 
 
+@pytest.mark.parametrize("centered", [-0.3, 10.0, np.array([0.1, -2.0])])
+def test_loc_scale_centered_invalid(centered):
+    N = 10
+    loc = np.random.uniform(-1.0, 1.0, size=(N, 2))
+    scale = np.random.uniform(0.5, 1.5, size=(N, 2))
+
+    def model(loc, scale):
+        with numpyro.plate("particles", N):
+            numpyro.sample("x", dist.Normal(loc, scale).to_event(1))
+
+    with pytest.raises(ValueError, match=".*does not satisfy.*"):
+        with handlers.reparam(config=LocScaleReparam(centered)):
+            with handlers.seed(rng_seed=10):
+                model(loc, scale)
+
+
 @pytest.mark.parametrize("shape", [(), (4,), (3, 2)], ids=str)
 @pytest.mark.parametrize("dim", [2, 3, 4])
 def test_projected_normal(shape, dim):