Repository: pyro-ppl/numpyro

Issue:
Issue with `AutoContinuous` guides when using `numpyro.deterministic`
I've been having some issues fitting guides that inherit from `AutoContinuous` (first noticed with `AutoDiagonalNormal`) to a model that has both parameter and deterministic sites. This problem does not occur for `AutoNormal`. I've included a toy example below, but for context I originally encountered this when trying to train a neural network with Bayesian last layer and wanted to use `numpyro.deterministic` to capture the logits I was passing to a categorical sampling statement.

Here's the example

```python
import numpyro
import numpyro.distributions as dist
from jax import random
from numpyro.distributions import constraints
from numpyro.infer import SVI, Trace_ELBO
from numpyro.infer.autoguide import AutoDiagonalNormal


def model(y=None):
    n = y.size if y is not None else 100

    mu = numpyro.sample("mu", dist.Normal(0, 5))
    sigma = numpyro.param("sigma", 1, constraint=constraints.positive)
    numpyro.deterministic("sigma2", sigma ** 2)

    numpyro.sample("y", dist.Normal(mu, sigma).expand((n,)), obs=y)


if __name__ == "__main__":
    rng_key = random.PRNGKey(0)
    mu, sigma = 2, 3

    rng_key, subkey = random.split(rng_key)
    y = mu + sigma * random.normal(subkey, shape=(1000,))

    guide = AutoDiagonalNormal(model)
    optimiser = numpyro.optim.Adam(step_size=0.01)
    svi = SVI(model, guide, optimiser, Trace_ELBO())

    rng_key, fit_key, sample_key = random.split(rng_key, 3)
    params, losses = svi.run(fit_key, num_steps=400, y=y)

    guide.sample_posterior(sample_key, params, sample_shape=(1000,))
```

if you run this you get the error

```
AttributeError: 'function' object has no attribute 'support'
```

I believe I've tracked the problem down. The addition of the deterministic site means that `self._postprocess_fn` is built from `numpyro.infer.util.constrain_fn` rather than `numpyro.infer.util.transform_fn`, that happens [here](https://github.com/pyro-ppl/numpyro/blob/048d2c80d9f4087aa9614225568bb88e1f74d669/numpyro/infer/util.py#L401-L411). `numpyro.infer.util.constrain_fn` has [this](https://github.com/pyro-ppl/numpyro/blob/048d2c80d9f4087aa9614225568bb88e1f74d669/numpyro/infer/util.py#L107-L109) problematic substitution function which is where the error originates from, as it tries to access `site["fn"].support`, but this doesn't exist for parameter sites. The parameter sites are added [here](https://github.com/pyro-ppl/numpyro/blob/048d2c80d9f4087aa9614225568bb88e1f74d669/numpyro/infer/autoguide.py#L505-L513) because they are needed to replay the model.

Would a valid fix be to check if the site corresponds to a parameter [here](https://github.com/pyro-ppl/numpyro/blob/048d2c80d9f4087aa9614225568bb88e1f74d669/numpyro/infer/util.py#L108-L109) and simply return the value if so without trying to biject it? Something like

```python
def substitute_fn(site):
    if site["name"] in params:
        if site["type"] == "sample":
            return biject_to(site["fn"].support)(params[site["name"]])
        else:
            return params[site["name"]]
```

Would be happy to submit a PR if this looks right?


Discussion:
Good catch - I believe all reasoning is correct. Please submit a PR (with this regression test in `test/infer/test_autoguide.py`). Thanks for the detailed explaination, @tcbegley!

Resolving patch:
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -106,7 +106,10 @@ def constrain_fn(model, model_args, model_kwargs, params, return_deterministic=F
 
     def substitute_fn(site):
         if site["name"] in params:
-            return biject_to(site["fn"].support)(params[site["name"]])
+            if site["type"] == "sample":
+                return biject_to(site["fn"].support)(params[site["name"]])
+            else:
+                return params[site["name"]]
 
     substituted_model = substitute(model, substitute_fn=substitute_fn)
     model_trace = trace(substituted_model).get_trace(*model_args, **model_kwargs)


Test patch:
diff --git a/test/infer/test_autoguide.py b/test/infer/test_autoguide.py
--- a/test/infer/test_autoguide.py
+++ b/test/infer/test_autoguide.py
@@ -434,3 +434,51 @@ def create_plates(batch, subsample, full_size):
             batch = data[:, beg:end]
             beg = end
             svi_state, loss = update_fn(svi_state, batch, subsample, full_size)
+
+
+@pytest.mark.parametrize(
+    "auto_class",
+    [
+        AutoDiagonalNormal,
+        AutoMultivariateNormal,
+        AutoLaplaceApproximation,
+        AutoLowRankMultivariateNormal,
+        AutoNormal,
+        AutoDelta,
+    ],
+)
+def test_autoguide_deterministic(auto_class):
+    def model(y=None):
+        n = y.size if y is not None else 1
+
+        mu = numpyro.sample("mu", dist.Normal(0, 5))
+        sigma = numpyro.param("sigma", 1, constraint=constraints.positive)
+
+        y = numpyro.sample("y", dist.Normal(mu, sigma).expand((n,)), obs=y)
+        numpyro.deterministic("z", (y - mu) / sigma)
+
+    mu, sigma = 2, 3
+    y = mu + sigma * random.normal(random.PRNGKey(0), shape=(300,))
+    y_train = y[:200]
+    y_test = y[200:]
+
+    guide = auto_class(model)
+    optimiser = numpyro.optim.Adam(step_size=0.01)
+    svi = SVI(model, guide, optimiser, Trace_ELBO())
+
+    params, losses = svi.run(random.PRNGKey(0), num_steps=500, y=y_train)
+    posterior_samples = guide.sample_posterior(
+        random.PRNGKey(0), params, sample_shape=(1000,)
+    )
+
+    predictive = Predictive(model, posterior_samples, params=params)
+    predictive_samples = predictive(random.PRNGKey(0), y_test)
+
+    assert predictive_samples["y"].shape == (1000, 100)
+    assert predictive_samples["z"].shape == (1000, 100)
+    assert_allclose(
+        (predictive_samples["y"] - posterior_samples["mu"][..., None])
+        / params["sigma"],
+        predictive_samples["z"],
+        atol=0.05,
+    )