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:
Sorry for the breakage! Could you try to use the dev branch of lightweight mmm? I will ping a dev there for a release if it works.
I think it's related to numpyro. The problem function is numpyro.deterministic. 
Everything else works. 
I'll have a look but I reckon it's related to the meridian release
Do you mean that `pip install --upgrade git+https://github.com/google/lightweight_mmm.git` does not resolve the issue?
@fehiepsi saw your fix on lightweight `Change-Id: I7c0658b0a13506c319fd3e6e00cdf2791d64e26f`.

I believe the long-term fix here is 2-fold:
  1. Return deterministic sites in posterior_samples (mcmc saves deterministic sites in its samples, and accessed via mcmc.get_samples()).
  2. Predictive always pops deterministic sites.
  
If these are unfeasible for deeper reasons, then at least mention the pop trick here: https://num.pyro.ai/en/v0.2.0/utilities.html

As the current behavior is a bit counterintuitive.
I'm running into the same issue, here's a reproducible example:

```python
import numpy as np
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS,Predictive
from jax import random


X = np.random.normal(0, 1, size=1000)
y = 5 + 1.2*X + np.random.normal(size=1000)

def model(X,y=None):
    alpha = numpyro.sample("alpha", dist.Normal(0,10))
    beta = numpyro.sample("beta", dist.Normal(0,1))
    sigma = numpyro.sample("sigma", dist.Exponential(1))
    with numpyro.plate("data", len(X)):
        eta = numpyro.deterministic("eta", alpha + beta*X)
        obs = numpyro.sample("obs", dist.Normal(eta, sigma), obs=y)
   
# Run NUTS.
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000)
mcmc.run(random.PRNGKey(0), X=X, y=y)

# Make predictions where X is a different shape
posterior_samples = mcmc.get_samples()
# posterior_samples.pop("eta") # this fixes the issues
pred_func = Predictive(model, posterior_samples=posterior_samples)
```

<details><summary>traceback</summary>
<p>

```python
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
    [... skipping hidden 1 frame]

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/jax/_src/util.py:290, in cache.<locals>.wrap.<locals>.wrapper(*args, **kwargs)
    289 else:
--> 290   return cached(config.trace_context(), *args, **kwargs)

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/jax/_src/util.py:283, in cache.<locals>.wrap.<locals>.cached(_, *args, **kwargs)
    281 @functools.lru_cache(max_size)
    282 def cached(_, *args, **kwargs):
--> 283   return f(*args, **kwargs)

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/jax/_src/lax/lax.py:155, in _broadcast_shapes_cached(*shapes)
    153 @cache()
    154 def _broadcast_shapes_cached(*shapes: tuple[int, ...]) -> tuple[int, ...]:
--> 155   return _broadcast_shapes_uncached(*shapes)

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/jax/_src/lax/lax.py:171, in _broadcast_shapes_uncached(*shapes)
    170 if result_shape is None:
--> 171   raise ValueError(f"Incompatible shapes for broadcasting: shapes={list(shapes)}")
    172 return result_shape

ValueError: Incompatible shapes for broadcasting: shapes=[(200,), (1000,)]

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
Cell In[1], line 26
     24 # Make predictions where X is a different shape
     25 pred_func = Predictive(model, posterior_samples=mcmc.get_samples())
---> 26 preds = pred_func(random.PRNGKey(1), X=X[:200], y=None)

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/infer/util.py:1011, in Predictive.__call__(self, rng_key, *args, **kwargs)
   1001 """
   1002 Returns dict of samples from the predictive distribution. By default, only sample sites not
   1003 contained in `posterior_samples` are returned. This can be modified by changing the
   (...)
   1008 :param kwargs: model kwargs.
   1009 """
   1010 if self.batch_ndims == 0 or self.params == {} or self.guide is None:
-> 1011     return self._call_with_params(rng_key, self.params, args, kwargs)
   1012 elif self.batch_ndims == 1:  # batch over parameters
   1013     batch_size = jnp.shape(tree_flatten(self.params)[0][0])[0]

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/infer/util.py:988, in Predictive._call_with_params(self, rng_key, params, args, kwargs)
    977     posterior_samples = _predictive(
    978         guide_rng_key,
    979         guide,
   (...)
    985         model_kwargs=kwargs,
    986     )
    987 model = substitute(self.model, self.params)
--> 988 return _predictive(
    989     rng_key,
    990     model,
    991     posterior_samples,
    992     self._batch_shape,
    993     return_sites=self.return_sites,
    994     infer_discrete=self.infer_discrete,
    995     parallel=self.parallel,
    996     model_args=args,
    997     model_kwargs=kwargs,
    998 )

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/infer/util.py:825, in _predictive(rng_key, model, posterior_samples, batch_shape, return_sites, infer_discrete, parallel, model_args, model_kwargs)
    823 rng_key = rng_key.reshape(batch_shape + key_shape)
    824 chunk_size = num_samples if parallel else 1
--> 825 return soft_vmap(
    826     single_prediction, (rng_key, posterior_samples), len(batch_shape), chunk_size
    827 )

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/util.py:419, in soft_vmap(fn, xs, batch_ndims, chunk_size)
    413     xs = tree_map(
    414         lambda x: jnp.reshape(x, prepend_shape + (chunk_size,) + jnp.shape(x)[1:]),
    415         xs,
    416     )
    417     fn = vmap(fn)
--> 419 ys = lax.map(fn, xs) if num_chunks > 1 else fn(xs)
    420 map_ndims = int(num_chunks > 1) + int(chunk_size > 1)
    421 ys = tree_map(
    422     lambda y: jnp.reshape(
    423         y, (int(np.prod(jnp.shape(y)[:map_ndims])),) + jnp.shape(y)[map_ndims:]
    424     )[:batch_size],
    425     ys,
    426 )

    [... skipping hidden 12 frame]

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/infer/util.py:798, in _predictive.<locals>.single_prediction(val)
    789     pred_samples = _sample_posterior(
    790         config_enumerate(condition(model, samples)),
    791         first_available_dim,
   (...)
    795         **model_kwargs,
    796     )
    797 else:
--> 798     model_trace = trace(
    799         seed(substitute(masked_model, samples), rng_key)
    800     ).get_trace(*model_args, **model_kwargs)
    801     pred_samples = {name: site["value"] for name, site in model_trace.items()}
    803 if return_sites is not None:

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/handlers.py:171, in trace.get_trace(self, *args, **kwargs)
    163 def get_trace(self, *args, **kwargs):
    164     """
    165     Run the wrapped callable and return the recorded trace.
    166 
   (...)
    169     :return: `OrderedDict` containing the execution trace.
    170     """
--> 171     self(*args, **kwargs)
    172     return self.trace

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/primitives.py:105, in Messenger.__call__(self, *args, **kwargs)
    103     return self
    104 with self:
--> 105     return self.fn(*args, **kwargs)

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/primitives.py:105, in Messenger.__call__(self, *args, **kwargs)
    103     return self
    104 with self:
--> 105     return self.fn(*args, **kwargs)

    [... skipping similar frames: Messenger.__call__ at line 105 (2 times)]

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/primitives.py:105, in Messenger.__call__(self, *args, **kwargs)
    103     return self
    104 with self:
--> 105     return self.fn(*args, **kwargs)

Cell In[1], line 17, in model(X, y)
     15 with numpyro.plate("data", len(X)):
     16     eta = numpyro.deterministic("eta", alpha + beta*X)
---> 17     obs = numpyro.sample("obs", dist.Normal(eta, sigma), obs=y)

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/primitives.py:222, in sample(name, fn, obs, rng_key, sample_shape, infer, obs_mask)
    207 initial_msg = {
    208     "type": "sample",
    209     "name": name,
   (...)
    218     "infer": {} if infer is None else infer,
    219 }
    221 # ...and use apply_stack to send it to the Messengers
--> 222 msg = apply_stack(initial_msg)
    223 return msg["value"]

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/primitives.py:47, in apply_stack(msg)
     45 pointer = 0
     46 for pointer, handler in enumerate(reversed(_PYRO_STACK)):
---> 47     handler.process_message(msg)
     48     # When a Messenger sets the "stop" field of a message,
     49     # it prevents any Messengers above it on the stack from being applied.
     50     if msg.get("stop"):

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/numpyro/primitives.py:546, in plate.process_message(self, msg)
    544 overlap_idx = max(len(expected_shape) - len(dist_batch_shape), 0)
    545 trailing_shape = expected_shape[overlap_idx:]
--> 546 broadcast_shape = lax.broadcast_shapes(
    547     trailing_shape, tuple(dist_batch_shape)
    548 )
    549 batch_shape = expected_shape[:overlap_idx] + broadcast_shape
    550 msg["fn"] = msg["fn"].expand(batch_shape)

    [... skipping hidden 1 frame]

File ~/Library/Caches/pypoetry/virtualenvs/ao-models-QNM_MrRk-py3.11/lib/python3.11/site-packages/jax/_src/lax/lax.py:171, in _broadcast_shapes_uncached(*shapes)
    169 result_shape = _try_broadcast_shapes(shape_list)
    170 if result_shape is None:
--> 171   raise ValueError(f"Incompatible shapes for broadcasting: shapes={list(shapes)}")
    172 return result_shape

ValueError: Incompatible shapes for broadcasting: shapes=[(200,), (1000,)]
```

</p>
</details>

I get that inputting samples for a deterministic site would lead to the model expecting a certain shape, but it does seem a bit awkward that the typical workflow with predictions requires some extra work if deterministics are involved. 

I wonder if something like this is possible? https://github.com/pyro-ppl/numpyro/blob/2f1bccdba2fc7b0a6ec235ca1bd5ce2417a0635c/numpyro/infer/mcmc.py#L714C61-L714C62
Hi @nikisix and @kylejcaron, really sorry for the breakage! I think a good action is to introduce `exclude_deterministic=True` to Predictive. This rolls the behavior back to pre-0.14 release. I'm less worried that new users will want to use deterministic sites in Predictive. What do you think, @martinjankowiak?
something like that sounds reasonable. the change in behavior was probably a mistake...
@fehiepsi @martinjankowiak should the `AutoGuide.sample_posterior()` be changed as well? It seems more difficult to fix since many `sample_posterior` functions are unique to auto guides. 

For example, the following workflow has the same problem :

```python
guide = AutoNormal(model) 
svi = SVI(model, guide, optim=numpyro.optim.Adam(0.01), loss=Trace_ELBO())
svi_result = svi.run(random.PRNGKey(0), 10000, X=X, y=y)

params = guide.sample_posterior(random.PRNGKey(0), params=svi_result.params)
pred_func = Predictive(model, params=params, num_samples=100)
preds = pred_func(random.PRNGKey(1), X=X[:250], y=None)
```

The solution for this seems to just including the guide and using SVI params instead, but I imagine some may be using the pattern above

```python
pred_func = Predictive(model, guide=guide, params=svi_result.params, num_samples=100)
preds = pred_func(random.PRNGKey(1),X[:n_preds])['eta']
```
I think [this pattern](https://github.com/pyro-ppl/numpyro/blob/2f1bccdba2fc7b0a6ec235ca1bd5ce2417a0635c/numpyro/infer/autoguide.py#L589C1-L590C1) could be used with an `exclude_deterministic` arg in AutoGuide's
@kylejcaron I think we can fix this in Predictive. The breakage happens because we allow substituting deterministic sites in the substitute handler. We can create a subclass of `substitute` which skip processing deterministic sites and use it in Predictive.
> @kylejcaron I think we can fix this in Predictive. The breakage happens because we allow substituting deterministic sites in the substitute handler. We can create a subclass of `substitute` which skip processing deterministic sites and use it in Predictive.

Got it that makes sense to me - seems like it'd involve just replacing the substitute call in [this line](https://github.com/pyro-ppl/numpyro/blob/d6ceae1bba0ccd76e7b68ae4c56ed911f5559096/numpyro/infer/util.py#L976C21-L976C31) and L987, but let me know if I'm missing anything.

I'm happy to make an attempt at this, any name recommendations for the new effect handler?

Resolving patch:
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- 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 single_prediction(val):
                 **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 @@ def __init__(
         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 @@ def __init__(
         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 @@ def _call_with_params(self, rng_key, params, args, kwargs):
                 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 @@ def _call_with_params(self, rng_key, params, args, kwargs):
             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
--- a/test/infer/test_infer_util.py
+++ b/test/infer/test_infer_util.py
@@ -53,6 +53,22 @@ def model(data=None):
     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)