Repository: pyro-ppl/numpyro

Issue:
numpyro.deterministic static on infer.Predictive
For some reason after fitting the model the numpyro.deterministic shape remains static, after trying to predict with a different shape it throws a shape error.

Example in lightweight-mmm:

```
# extra_features.shape = (10,3) / trying to predict 10 new time periods

extra_features_effect = numpyro.deterministic(
    name="extra_features_effect",
    value=jnp.einsum(
        extra_features_einsum, extra_features, coef_extra_features
    ),
)

# extra_features_effect.shape = (30,3)  / output is resized to the size of the model as when fit; 30 periods
```

This throws a size error, see:
https://github.com/google/lightweight_mmm/issues/309
and
https://github.com/google/lightweight_mmm/issues/308


Discussion:

Patch:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
index 8ffd3887..9b159187 100644
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -188,11 +188,23 @@ def _sample_fn_nojit_args(state, sampler, args, kwargs):
     return (sampler.sample(state[0], args, kwargs),)
 
 
-def _collect_fn(collect_fields):
-    @cached_by(_collect_fn, collect_fields)
+def _collect_fn(collect_fields, remove_sites):
+    @cached_by(_collect_fn, collect_fields, remove_sites)
     def collect(x):
         if collect_fields:
-            return attrgetter(*collect_fields)(x[0])
+            fields = attrgetter(*collect_fields)(x[0])
+
+            if remove_sites != ():
+                fields = [fields] if len(collect_fields) == 1 else list(fields)
+                assert isinstance(fields[0], dict)
+
+                sample_sites = fields[0].copy()
+                for site in remove_sites:
+                    sample_sites.pop(site)
+                fields[0] = sample_sites
+                fields = fields[0] if len(collect_fields) == 1 else fields
+
+            return fields
         else:
             return x[0]
 
@@ -419,7 +431,7 @@ class MCMC(object):
         except TypeError:
             return None
 
-    def _single_chain_mcmc(self, init, args, kwargs, collect_fields):
+    def _single_chain_mcmc(self, init, args, kwargs, collect_fields, remove_sites):
         rng_key, init_state, init_params = init
         # Check if _sample_fn is None, then we need to initialize the sampler.
         if init_state is None or (getattr(self.sampler, "_sample_fn", None) is None):
@@ -452,7 +464,7 @@ class MCMC(object):
             upper_idx,
             sample_fn,
             init_val,
-            transform=_collect_fn(collect_fields),
+            transform=_collect_fn(collect_fields, remove_sites),
             progbar=self.progress_bar,
             return_last_val=True,
             thinning=self.thinning,
@@ -556,7 +568,8 @@ class MCMC(object):
             These are typically the arguments needed by the `model`.
         :param extra_fields: Extra fields (aside from :meth:`~numpyro.infer.mcmc.MCMCKernel.default_fields`)
             from the state object (e.g. :data:`numpyro.infer.hmc.HMCState` for HMC) to collect during
-            the MCMC run.
+            the MCMC run. Exclude sample sites from collection with "~`sampler.sample_field`.`sample_site`".
+            e.g. "~z.a" will prevent site "a" from being collected if you're using the NUTS sampler.
         :type extra_fields: tuple or list
         :param bool collect_warmup: Whether to collect samples from the warmup phase. Defaults
             to `False`.
@@ -591,7 +604,9 @@ class MCMC(object):
         :param extra_fields: Extra fields (aside from `"z"`, `"diverging"`) from the
             state object (e.g. :data:`numpyro.infer.hmc.HMCState` for HMC) to be collected
             during the MCMC run. Note that subfields can be accessed using dots, e.g.
-            `"adapt_state.step_size"` can be used to collect step sizes at each step.
+            `"adapt_state.step_size"` can be used to collect step sizes at each step. Exclude sample sites from
+            collection with "~`sampler.sample_field`.`sample_site`". e.g. "~z.a" will prevent site "a" from
+            being collected if you're using the NUTS sampler.
         :type extra_fields: tuple or list of str
         :param init_params: Initial parameters to begin sampling. The type must be consistent
             with the input type to `potential_fn` provided to the kernel. If the kernel is
@@ -626,18 +641,25 @@ class MCMC(object):
                     " as `num_chains`."
                 )
         assert isinstance(extra_fields, (tuple, list))
-        collect_fields = tuple(
-            set(
-                (self._sample_field,)
-                + tuple(self._default_fields)
-                + tuple(extra_fields)
-            )
-        )
+
+        collect_fields = {}
+        remove_sites = {}
+        for field_name in (
+            (self._sample_field,) + tuple(self._default_fields) + tuple(extra_fields)
+        ):
+            if field_name.startswith(f"~{self._sample_field}."):
+                remove_sites[(field_name[len(self._sample_field) + 2 :])] = None
+            else:
+                collect_fields[field_name] = None
+        collect_fields = tuple(collect_fields.keys())
+        remove_sites = tuple(remove_sites.keys())
+
         partial_map_fn = partial(
             self._single_chain_mcmc,
             args=args,
             kwargs=kwargs,
             collect_fields=collect_fields,
+            remove_sites=remove_sites,
         )
         map_args = (rng_key, init_state, init_params)
         if self.num_chains == 1:
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
index 916545cd..7d57cbbd 100644
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -761,6 +761,7 @@ def _predictive(
     return_sites=None,
     infer_discrete=False,
     parallel=True,
+    exclude_deterministic: bool = True,
     model_args=(),
     model_kwargs={},
 ):
@@ -774,7 +775,7 @@ def _predictive(
             posterior_samples,
         )
         prototype_trace = trace(
-            seed(substitute(masked_model, prototype_sample), subkey)
+            seed(condition(masked_model, prototype_sample), subkey)
         ).get_trace(*model_args, **model_kwargs)
         first_available_dim = -_guess_max_plate_nesting(prototype_trace) - 1
 
@@ -795,9 +796,20 @@ def _predictive(
                 **model_kwargs,
             )
         else:
-            model_trace = trace(
-                seed(substitute(masked_model, samples), rng_key)
-            ).get_trace(*model_args, **model_kwargs)
+
+            def _samples_wo_deterministic(msg):
+                return (
+                    samples.get(msg["name"]) if msg["type"] != "deterministic" else None
+                )
+
+            substituted_model = (
+                substitute(masked_model, substitute_fn=_samples_wo_deterministic)
+                if exclude_deterministic
+                else substitute(masked_model, samples)
+            )
+            model_trace = trace(seed(substituted_model, rng_key)).get_trace(
+                *model_args, **model_kwargs
+            )
             pred_samples = {name: site["value"] for name, site in model_trace.items()}
 
         if return_sites is not None:
@@ -870,6 +882,7 @@ class Predictive(object):
 
         + set `batch_ndims=1` to get predictions from a one dimensional batch of the guide and parameters
           with shapes `(num_samples x batch_size x ...)`
+    :param exclude_deterministic: indicates whether to ignore deterministic sites from the posterior samples.
 
     :return: dict of samples from the predictive distribution.
 
@@ -907,6 +920,7 @@ class Predictive(object):
         infer_discrete: bool = False,
         parallel: bool = False,
         batch_ndims: Optional[int] = None,
+        exclude_deterministic: bool = True,
     ):
         if posterior_samples is None and num_samples is None:
             raise ValueError(
@@ -967,6 +981,7 @@ class Predictive(object):
         self.parallel = parallel
         self.batch_ndims = batch_ndims
         self._batch_shape = batch_shape
+        self.exclude_deterministic = exclude_deterministic
 
     def _call_with_params(self, rng_key, params, args, kwargs):
         posterior_samples = self.posterior_samples
@@ -983,6 +998,7 @@ class Predictive(object):
                 parallel=self.parallel,
                 model_args=args,
                 model_kwargs=kwargs,
+                exclude_deterministic=self.exclude_deterministic,
             )
         model = substitute(self.model, self.params)
         return _predictive(
@@ -995,6 +1011,7 @@ class Predictive(object):
             parallel=self.parallel,
             model_args=args,
             model_kwargs=kwargs,
+            exclude_deterministic=self.exclude_deterministic,
         )
 
     def __call__(self, rng_key, *args, **kwargs):


Test patch:
diff --git a/test/infer/test_infer_util.py b/test/infer/test_infer_util.py
index 4412916b..aecf7ca2 100644
--- a/test/infer/test_infer_util.py
+++ b/test/infer/test_infer_util.py
@@ -53,6 +53,22 @@ def beta_bernoulli():
     return model, data, true_probs
 
 
+def linear_regression():
+    N = 800
+    X = dist.Normal(0, 1).sample(random.PRNGKey(0), (N,))
+    y = 1.5 + X * 0.7
+
+    def model(X, y=None):
+        alpha = numpyro.sample("alpha", dist.Normal(0.0, 5))
+        beta = numpyro.sample("beta", dist.Normal(0.0, 1.0))
+        sigma = numpyro.sample("sigma", dist.Exponential(1.0))
+        with numpyro.plate("plate", len(X)):
+            mu = numpyro.deterministic("mu", alpha + X * beta)
+            numpyro.sample("obs", dist.Normal(mu, sigma), obs=y)
+
+    return model, X, y
+
+
 @pytest.mark.parametrize("parallel", [True, False])
 def test_predictive(parallel):
     model, data, true_probs = beta_bernoulli()
@@ -74,6 +90,29 @@ def test_predictive(parallel):
     assert_allclose(obs.mean(0), true_probs, rtol=0.1)
 
 
+@pytest.mark.parametrize("parallel", [True, False])
+def test_predictive_with_deterministic(parallel):
+    """Tests that the default behavior when predicting from models with
+    deterministic sites doesn't lead to static deterministic sites in the predictive.
+    """
+    n_preds = 400
+    model, X, y = linear_regression()
+    mcmc = MCMC(NUTS(model), num_warmup=100, num_samples=100)
+    mcmc.run(random.PRNGKey(0), X=X, y=y)
+    samples = mcmc.get_samples()
+    predictive = Predictive(model, samples, parallel=parallel)
+    # change the input (X) shape to make sure the deterministic shape changes
+    predictive_samples = predictive(random.PRNGKey(1), X=X[:n_preds])
+    assert predictive_samples.keys() == {"mu", "obs"}
+
+    predictive.return_sites = ["beta", "mu", "obs"]
+    # change the input (X) shape to make sure the deterministic shape changes
+    predictive_samples = predictive(random.PRNGKey(1), X=X[:n_preds])
+    # check shapes
+    assert predictive_samples["mu"].shape == (100,) + X[:n_preds].shape
+    assert predictive_samples["obs"].shape == (100,) + X[:n_preds].shape
+
+
 def test_predictive_with_guide():
     data = jnp.array([1] * 8 + [0] * 2)
 
diff --git a/test/infer/test_mcmc.py b/test/infer/test_mcmc.py
index fa1fea46..480aba80 100644
--- a/test/infer/test_mcmc.py
+++ b/test/infer/test_mcmc.py
@@ -1186,3 +1186,17 @@ def test_none_trajectory_length(num_steps):
         mcmc.run(rng_key, data, extra_fields=("num_steps",))
         num_steps_list = np.array(mcmc.get_extra_fields()["num_steps"])
         assert all(step == num_steps for step in num_steps_list)
+
+
+@pytest.mark.parametrize("kernel_cls", [NUTS, BarkerMH])
+@pytest.mark.parametrize("remove_sites", [("~z.a", "~z.b"), ("~z.a", "~z.a")])
+def test_remove_sites(kernel_cls, remove_sites):
+    def model():
+        numpyro.sample("a", dist.Normal())
+        numpyro.sample("b", dist.Normal())
+
+    mcmc = MCMC(kernel_cls(model), num_warmup=10, num_samples=10)
+    mcmc.run(random.PRNGKey(0), extra_fields=remove_sites)
+    samps = mcmc.get_samples()
+
+    assert all([site[3:] not in samps for site in remove_sites])