Repository: pyro-ppl/numpyro

Issue:
LocScaleReparam samples outside the support
Hi 

I am a numpyro user having trouble with `LocScaleReparam ` function with non-normal distributions that still have loc/scale parameters. 

When I run LocScaleReparam with non-normal distributions, the reparametrized samples are outside the support. For instance, when I sample from a decentered LogNormal distribution (positive support), I get samples that are negative. 


```python
def model(dim=10):
    y = numpyro.sample("y", dist.LogNormal(0, 0.5))
    numpyro.sample("x", dist.LogNormal(-10*jnp.ones(dim - 1), y))
``` 

Using the code from [neal's funnel](https://num.pyro.ai/en/stable/examples/funnel.html), and printing outputs
```python
reparam_model = reparam(model, config={"x": LocScaleReparam(0)})

def run_inference(model, args, rng_key):
    kernel = NUTS(model)
    mcmc = MCMC(
        kernel,
        num_warmup=args.num_warmup,
        num_samples=args.num_samples,
        num_chains=args.num_chains,
        progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
    )
    mcmc.run(rng_key)
    mcmc.print_summary(exclude_deterministic=False)
    return mcmc


def main(args):
    rng_key = random.PRNGKey(0)

    # do inference with centered parameterization
    print(
        "============================= Centered Parameterization =============================="
    )
    mcmc = run_inference(model, args, rng_key)
    samples = mcmc.get_samples()
    diverging = mcmc.get_extra_fields()["diverging"]

    # do inference with non-centered parameterization
    print(
        "\n=========================== Non-centered Parameterization ============================"
    )
    reparam_mcmc = run_inference(reparam_model, args, rng_key)
    reparam_samples = reparam_mcmc.get_samples()
    reparam_diverging = reparam_mcmc.get_extra_fields()["diverging"]
    # collect deterministic sites
    reparam_samples = Predictive(
        reparam_model, reparam_samples, return_sites=["x", "y"]
    )(random.PRNGKey(1))

    print(f"min sample: {jnp.min(samples['x'])}")
    print(f"reparametrized min sample: {jnp.min(reparam_samples['x'])}")

if __name__ == "__main__":
    assert numpyro.__version__.startswith("0.10.1")
    parser = [argparse.ArgumentParser](https://docs.python.org/3/library/argparse.html#argparse.ArgumentParser)(
        description="Non-centered reparameterization example"
    )
    parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
    parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
    parser.add_argument("--num-chains", nargs="?", default=1, type=int)
    parser.add_argument("--device", default="cpu", type=str, help='use "cpu" or "gpu".')
    args = parser.parse_args()

    numpyro.set_platform(args.device)
    numpyro.set_host_device_count(args.num_chains)

    main(args)

```
output:
```
min sample: 1.0087332746877564e-08
reparametrized min sample: -9.989072799682617
```

I am guessing that either I have implemented this wrong or 
the reparametrization is happening in the original space, and not the transformed space for the TransformedDistributions (e.g. reparam.py L129-130)

Thank you for your help!


Discussion:
Hi @joshryu0425, LocScaleReparam cannot be used for positive support. We will raise error for this. You can use this reparam for real-support distributions like Cauchy, StudentT,...

Resolving patch:
diff --git a/numpyro/infer/reparam.py b/numpyro/infer/reparam.py
--- a/numpyro/infer/reparam.py
+++ b/numpyro/infer/reparam.py
@@ -100,6 +100,15 @@ def __init__(self, centered=None, shape_params=()):
 
     def __call__(self, name, fn, obs):
         assert obs is None, "LocScaleReparam does not support observe statements"
+        support = fn.support
+        if isinstance(support, constraints.independent):
+            support = fn.support.base_constraint
+        if support is not constraints.real:
+            raise ValueError(
+                "LocScaleReparam only supports distributions with real "
+                f"support, but got {support} support at site {name}."
+            )
+
         centered = self.centered
         if is_identically_one(centered):
             return fn, obs


Test patch:
diff --git a/test/infer/test_reparam.py b/test/infer/test_reparam.py
--- a/test/infer/test_reparam.py
+++ b/test/infer/test_reparam.py
@@ -390,3 +390,12 @@ def model():
     )
     log_det_site = handlers.trace(substituted_model).get_trace()["_x_log_det"]
     assert log_det_site["scale"] == 5.0
+
+
+def test_loc_scale_reparam_raise_for_log_normal():
+    def model():
+        numpyro.sample("x", dist.LogNormal(0, 1))
+
+    reparam_model = handlers.reparam(model, config={"x": LocScaleReparam(0)})
+    with pytest.raises(ValueError, match="LocScaleReparam.*"):
+        handlers.seed(reparam_model, rng_seed=0)()