Repository: pyro-ppl/numpyro

Issue:
raise a better error/warning for init strategies with substituted models
Right now, init strategy with a ```substituted``` model breaks, and the error message showed is not straight forward. It could be clarified that, in order to work, ```substitute``` requires that the substituted site is ```blocked```, and that anyway ```condition``` is preferred to ```substitute```  for modelling purposes.


Discussion:

Resolving patch:
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -112,7 +112,7 @@ class trace(Messenger):
     Returns a handler that records the inputs and outputs at primitive calls
     inside `fn`.
 
-    **Example**
+    **Example:**
 
     .. doctest::
 
@@ -175,7 +175,7 @@ class replay(Messenger):
     :param fn: Python callable with NumPyro primitives.
     :param guide_trace: an OrderedDict containing execution metadata.
 
-    **Example**
+    **Example:**
 
     .. doctest::
 
@@ -600,14 +600,14 @@ class scope(Messenger):
     """
     This handler prepend a prefix followed by a divider to the name of sample sites.
 
-    **Example**
+    **Example:**
 
     .. doctest::
 
         >>> import numpyro
         >>> import numpyro.distributions as dist
         >>> from numpyro.handlers import scope, seed, trace
-        >>>
+
         >>> def model():
         ...     with scope(prefix="a"):
         ...         with scope(prefix="b", divider="."):
@@ -711,6 +711,10 @@ class substitute(Messenger):
     replaced by the value returned from the call to `substitute_fn`
     for the given site.
 
+    .. note:: This handler is mainly used for internal algorithms.
+        For conditioning a generative model on observed data, please
+        using the :class:`condition` handler.
+
     :param fn: Python callable with NumPyro primitives.
     :param dict data: dictionary of `numpy.ndarray` values keyed by
         site names.
diff --git a/numpyro/infer/initialization.py b/numpyro/infer/initialization.py
--- a/numpyro/infer/initialization.py
+++ b/numpyro/infer/initialization.py
@@ -2,8 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from functools import partial
+import warnings
 
-from jax import random
 import jax.numpy as jnp
 
 import numpyro.distributions as dist
@@ -25,6 +25,13 @@ def init_to_median(site=None, num_samples=15):
         and not site["is_observed"]
         and not site["fn"].is_discrete
     ):
+        if site["value"] is not None:
+            warnings.warn(
+                f"init_to_median() skipping initialization of site '{site['name']}'"
+                " which already stores a value."
+            )
+            return site["value"]
+
         rng_key = site["kwargs"].get("rng_key")
         sample_shape = site["kwargs"].get("sample_shape")
         try:
@@ -58,12 +65,18 @@ def init_to_uniform(site=None, radius=2):
         and not site["is_observed"]
         and not site["fn"].is_discrete
     ):
+        if site["value"] is not None:
+            warnings.warn(
+                f"init_to_uniform() skipping initialization of site '{site['name']}'"
+                " which already stores a value."
+            )
+            return site["value"]
+
         # XXX: we import here to avoid circular import
         from numpyro.infer.util import helpful_support_errors
 
         rng_key = site["kwargs"].get("rng_key")
         sample_shape = site["kwargs"].get("sample_shape")
-        rng_key, subkey = random.split(rng_key)
 
         with helpful_support_errors(site):
             transform = biject_to(site["fn"].support)


Test patch:
diff --git a/test/infer/test_mcmc.py b/test/infer/test_mcmc.py
--- a/test/infer/test_mcmc.py
+++ b/test/infer/test_mcmc.py
@@ -1049,3 +1049,14 @@ def model():
     assert set(inverse_mass_matrix.keys()) == {("x", "z")}
     expected_mm = jnp.diag(expected_mm) if dense_mass else expected_mm
     assert_allclose(inverse_mass_matrix[("x", "z")], expected_mm)
+
+
+def test_init_strategy_substituted_model():
+    def model():
+        numpyro.sample("x", dist.Normal(0, 1))
+        numpyro.sample("y", dist.Normal(0, 1))
+
+    subs_model = numpyro.handlers.substitute(model, data={"x": 10.0})
+    mcmc = MCMC(NUTS(subs_model), num_warmup=10, num_samples=10)
+    with pytest.warns(UserWarning, match="skipping initialization"):
+        mcmc.run(random.PRNGKey(1))