Repository: pyro-ppl/numpyro

Issue:
NumPyro doesn't raise if a determinstic variable has the same name as a random variable
E.g.:
```python
import numpyro
import numpyro.distributions as dist
from numpyro.infer import NUTS, MCMC
import jax


def model():
    alpha = numpyro.sample("alpha", dist.Normal(0, 1))
    numpyro.deterministic("alpha", alpha * 2)


mcmc = MCMC(NUTS(model), num_chains=1, num_samples=10, num_warmup=10)
mcmc.run(jax.random.PRNGKey(0))
```

If executed, this gives the rather mysterious
```
$ python t.py 
Traceback (most recent call last):
  File "t.py", line 13, in <module>
    mcmc.run(jax.random.PRNGKey(0))
  File "/home/marco/numpyro-dev/numpyro/infer/mcmc.py", line 599, in run
    states_flat, last_state = partial_map_fn(map_args)
  File "/home/marco/numpyro-dev/numpyro/infer/mcmc.py", line 387, in _single_chain_mcmc
    init_state = self.sampler.init(
  File "/home/marco/numpyro-dev/numpyro/infer/hmc.py", line 696, in init
    init_params = self._init_state(
  File "/home/marco/numpyro-dev/numpyro/infer/hmc.py", line 642, in _init_state
    init_params, potential_fn, postprocess_fn, model_trace = initialize_model(
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 653, in initialize_model
    (init_params, pe, grad), is_valid = find_valid_initial_params(
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 389, in find_valid_initial_params
    (init_params, pe, z_grad), is_valid = _find_valid_params(
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 375, in _find_valid_params
    _, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state)
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 360, in body_fn
    pe, z_grad = value_and_grad(potential_fn)(params)
  File "/home/marco/numpyro-dev/.venv/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/marco/numpyro-dev/.venv/lib/python3.8/site-packages/jax/_src/api.py", line 993, in value_and_grad_f
    ans, vjp_py = _vjp(f_partial, *dyn_args, reduce_axes=reduce_axes)
  File "/home/marco/numpyro-dev/.venv/lib/python3.8/site-packages/jax/_src/api.py", line 2312, in _vjp
    out_primal, out_vjp = ad.vjp(
  File "/home/marco/numpyro-dev/.venv/lib/python3.8/site-packages/jax/interpreters/ad.py", line 116, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  File "/home/marco/numpyro-dev/.venv/lib/python3.8/site-packages/jax/interpreters/ad.py", line 103, in linearize
    jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
  File "/home/marco/numpyro-dev/.venv/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 513, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  File "/home/marco/numpyro-dev/.venv/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 242, in potential_energy
    log_joint, model_trace = log_density_(
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 56, in log_density
    model_trace = trace(model).get_trace(*model_args, **model_kwargs)
  File "/home/marco/numpyro-dev/numpyro/handlers.py", line 171, in get_trace
    self(*args, **kwargs)
  File "/home/marco/numpyro-dev/numpyro/primitives.py", line 87, in __call__
    return self.fn(*args, **kwargs)
  File "/home/marco/numpyro-dev/numpyro/primitives.py", line 87, in __call__
    return self.fn(*args, **kwargs)
  File "/home/marco/numpyro-dev/numpyro/primitives.py", line 87, in __call__
    return self.fn(*args, **kwargs)
  [Previous line repeated 2 more times]
  File "t.py", line 8, in model
    alpha = numpyro.sample("alpha", dist.Normal(0, 1))
  File "/home/marco/numpyro-dev/numpyro/primitives.py", line 201, in sample
    msg = apply_stack(initial_msg)
  File "/home/marco/numpyro-dev/numpyro/primitives.py", line 31, in apply_stack
    msg["value"], msg["intermediates"] = msg["fn"](
  File "/home/marco/numpyro-dev/numpyro/distributions/distribution.py", line 303, in __call__
    return self.sample_with_intermediates(key, *args, **kwargs)
  File "/home/marco/numpyro-dev/numpyro/distributions/distribution.py", line 261, in sample_with_intermediates
    return self.sample(key, sample_shape=sample_shape), []
  File "/home/marco/numpyro-dev/numpyro/distributions/continuous.py", line 1358, in sample
    assert is_prng_key(key)
jax._src.traceback_util.UnfilteredStackTrace: AssertionError

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "t.py", line 13, in <module>
    mcmc.run(jax.random.PRNGKey(0))
  File "/home/marco/numpyro-dev/numpyro/infer/mcmc.py", line 599, in run
    states_flat, last_state = partial_map_fn(map_args)
  File "/home/marco/numpyro-dev/numpyro/infer/mcmc.py", line 387, in _single_chain_mcmc
    init_state = self.sampler.init(
  File "/home/marco/numpyro-dev/numpyro/infer/hmc.py", line 696, in init
    init_params = self._init_state(
  File "/home/marco/numpyro-dev/numpyro/infer/hmc.py", line 642, in _init_state
    init_params, potential_fn, postprocess_fn, model_trace = initialize_model(
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 653, in initialize_model
    (init_params, pe, grad), is_valid = find_valid_initial_params(
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 389, in find_valid_initial_params
    (init_params, pe, z_grad), is_valid = _find_valid_params(
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 375, in _find_valid_params
    _, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state)
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 360, in body_fn
    pe, z_grad = value_and_grad(potential_fn)(params)
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 242, in potential_energy
    log_joint, model_trace = log_density_(
  File "/home/marco/numpyro-dev/numpyro/infer/util.py", line 56, in log_density
    model_trace = trace(model).get_trace(*model_args, **model_kwargs)
  File "/home/marco/numpyro-dev/numpyro/handlers.py", line 171, in get_trace
    self(*args, **kwargs)
  File "/home/marco/numpyro-dev/numpyro/primitives.py", line 87, in __call__
    return self.fn(*args, **kwargs)
  File "/home/marco/numpyro-dev/numpyro/primitives.py", line 87, in __call__
    return self.fn(*args, **kwargs)
  File "/home/marco/numpyro-dev/numpyro/primitives.py", line 87, in __call__
    return self.fn(*args, **kwargs)
  [Previous line repeated 2 more times]
  File "t.py", line 8, in model
    alpha = numpyro.sample("alpha", dist.Normal(0, 1))
  File "/home/marco/numpyro-dev/numpyro/primitives.py", line 201, in sample
    msg = apply_stack(initial_msg)
  File "/home/marco/numpyro-dev/numpyro/primitives.py", line 31, in apply_stack
    msg["value"], msg["intermediates"] = msg["fn"](
  File "/home/marco/numpyro-dev/numpyro/distributions/distribution.py", line 303, in __call__
    return self.sample_with_intermediates(key, *args, **kwargs)
  File "/home/marco/numpyro-dev/numpyro/distributions/distribution.py", line 261, in sample_with_intermediates
    return self.sample(key, sample_shape=sample_shape), []
  File "/home/marco/numpyro-dev/numpyro/distributions/continuous.py", line 1358, in sample
    assert is_prng_key(key)
AssertionError
```



Discussion:

Resolving patch:
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -154,7 +154,7 @@ def postprocess_message(self, msg):
             # which has no name
             return
         assert not (
-            msg["type"] == "sample" and msg["name"] in self.trace
+            msg["type"] in ("sample", "deterministic") and msg["name"] in self.trace
         ), "all sites must have unique names but got `{}` duplicated".format(
             msg["name"]
         )


Test patch:
diff --git a/test/test_handlers.py b/test/test_handlers.py
--- a/test/test_handlers.py
+++ b/test/test_handlers.py
@@ -12,7 +12,7 @@
 from numpyro import handlers
 import numpyro.distributions as dist
 from numpyro.distributions import constraints
-from numpyro.infer import SVI, Trace_ELBO
+from numpyro.infer import MCMC, NUTS, SVI, Trace_ELBO
 from numpyro.infer.util import log_density
 import numpyro.optim as optim
 from numpyro.util import not_jax_tracer, optional
@@ -778,3 +778,14 @@ def subsample_fn(rng_key):
 
         # test that values are not duplicated
         assert len(set(subsamples[k].copy())) == subsample_size
+
+
+def test_sites_have_unique_names():
+    def model():
+        alpha = numpyro.sample("alpha", dist.Normal())
+        numpyro.deterministic("alpha", alpha * 2)
+
+    mcmc = MCMC(NUTS(model), num_chains=1, num_samples=10, num_warmup=10)
+    msg = "all sites must have unique names but got `alpha` duplicated"
+    with pytest.raises(AssertionError, match=msg):
+        mcmc.run(random.PRNGKey(0))