Predictive.exclude_deterministic does not filter deterministic sites

[brendancooley]

Bug Description

The Predictive helper class attribute exclude_deterministic does not successfully exclude deterministic sites from resulting posterior predictive dictionary.

Steps to Reproduce

uv run test_exclude_deterministic.py

where exclude_deterministic.py contains

#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.9"
# dependencies = [
#     "jax",
#     "jaxlib",
#     "numpyro",
# ]
# ///
"""
Test to verify if exclude_deterministic works in numpyro.infer.util.Predictive
"""

import jax
import jax.numpy as jnp

import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, Predictive


def model(x=None, y=None):
    """Simple model with both stochastic and deterministic sites"""
    # Stochastic sites
    a = numpyro.sample("a", dist.Normal(0, 1))
    b = numpyro.sample("b", dist.Normal(0, 1))

    # Deterministic site
    linear_combination = numpyro.deterministic("linear_combination", a + 2 * b)

    # Another deterministic site
    _ = numpyro.deterministic("squared", a**2)

    # Likelihood
    if y is not None:
        sigma = numpyro.sample("sigma", dist.Exponential(1.0))
        numpyro.sample("obs", dist.Normal(linear_combination, sigma), obs=y)


# Generate some synthetic data
key = jax.random.PRNGKey(0)
n_samples = 100
x_data = jnp.linspace(0, 1, n_samples)
true_a = 1.5
true_b = 2.0
y_data = true_a + 2 * true_b + jax.random.normal(key, (n_samples,)) * 0.5

# Run MCMC
mcmc = MCMC(NUTS(model), num_warmup=500, num_samples=1000)
mcmc.run(key, y=y_data)
samples = mcmc.get_samples()

# Test Predictive with exclude_deterministic=True
print("\nTest 1: Predictive with exclude_deterministic=True")
print("-" * 70)
predictive_exclude = Predictive(model, samples, exclude_deterministic=True)
predictions_exclude = predictive_exclude(key)
print("Keys with exclude_deterministic=True:", predictions_exclude.keys())
print("Expected: Should NOT include 'linear_combination' or 'squared'")

# Test Predictive with exclude_deterministic=False (default)
print("\nTest 2: Predictive with exclude_deterministic=False")
print("-" * 70)
predictive_include = Predictive(model, samples, exclude_deterministic=False)
predictions_include = predictive_include(key)
print("Keys with exclude_deterministic=False:", predictions_include.keys())
print("Expected: SHOULD include 'linear_combination' and 'squared'")

Expected Behavior

The "squared" and "linear_combination" sites should only be present in the returned value of predictive_include().

Notes

Obviously the fix here will be correct but breaking. Unless I'm missing something. Happy to work on a fix, but perhaps there's some context I'm missing.

[juanitorduz]

I agree this is confusing, from what I see from the PR adding this argument: https://github.com/pyro-ppl/numpyro/pull/1789/files. The idea is to exclude_deterministic in the substitute effect handler to handle the samples. This is especially important for new shapes. I think that the returned sites are supposed to be explicitly passed in the return_sites argument, see

numpyro/numpyro/infer/util.py

Lines 862 to 876 in 725e009

	if return_sites is not None:
	if return_sites == "":
	sites = {
	k for k, site in model_trace.items() if site["type"] != "plate"
	}
	else:
	sites = return_sites
	else:
	sites = {
	k
	for k, site in model_trace.items()
	if (site["type"] == "sample" and k not in samples)
	or (site["type"] == "deterministic")
	}
	return {name: value for name, value in pred_samples.items() if name in sites}

The confusion comes from the return_sites vs exclude_deterministic (I was also confused when trying to "fix" this bug)

[brendancooley]

I agree this is confusing, from what I see from the PR adding this argument: https://github.com/pyro-ppl/numpyro/pull/1789/files. The idea is to exclude_deterministic in the substitute effect handler to handle the samples. This is especially important for new shapes. I think that the returned sites are supposed to be explicitly passed in the return_sites argument, see

numpyro/numpyro/infer/util.py

Lines 862 to 876 in 725e009

	if return_sites is not None:
	if return_sites == "":
	sites = {
	k for k, site in model_trace.items() if site["type"] != "plate"
	}
	else:
	sites = return_sites
	else:
	sites = {
	k
	for k, site in model_trace.items()
	if (site["type"] == "sample" and k not in samples)
	or (site["type"] == "deterministic")
	}
	return {name: value for name, value in pred_samples.items() if name in sites}

The confusion comes from the return_sites vs exclude_deterministic (I was also confused when trying to "fix" this bug)

Ah, this is helpful. Thanks Juan. Maybe I'll keep this open and take a hack at improving the docs.

[kylejcaron]

Ah, this is helpful. Thanks Juan. Maybe I'll keep this open and take a hack at improving the docs.

As the person who made that change I also find it a bit confusing and would welcome new ideas or doc improvements.

The main reason behind that change was that a breakage occurred when Predictive was called - it would condition on the sites passed in, including deterministic sites. So if your deterministic site had size N=1000, but you want to predict on X_test with size N=500, it would break. This change made it the default behavior to re-sample deterministic sites from their parent nodes when making predictions as opposed to conditioning on them.

maybe changing the argument name to condition_deterministic: bool = False would help?

[juanitorduz]

I think this might be a better argument name indeed :)