Repository: pyro-ppl/numpyro

Issue:
Correct SVI API usage
So I was playing around with SVI API and problems from Statistical Rethinking 2ed. And one thing - extracting posterior sample data - made me quite puzzled for a while.

The [doc on SVI](https://num.pyro.ai/en/stable/svi.html#numpyro.infer.svi.SVI) describes intended usage as:
```python
# ...
# get posterior samples
predictive = Predictive(guide, params=params, num_samples=1000)
samples = predictive(random.PRNGKey(1), data)
```

The result of that for my model was: 
1. Sample with deterministic estimands missing if I used `Predictive(guide, params=params, num_samples=1000)`
```python
>>> samples.keys()
# actual: dict_keys(['alpha', 'beta', 'sigma'])
# expected: dict_keys(['W', 'alpha', 'beta', 'mu', 'sigma'])
```
2. Empty sample if I used `Predictive(guide, params=params, num_samples=1000)`
```python
>>> samples.keys()
# actual: dict_keys([])
# expected: dict_keys(['W', 'mu'])
```
---
Changing a snippet like this solved the issue:
```python
# ...
# get posterior samples
predictive = Predictive(guide.model, params=params, num_samples=1000)
samples = predictive(random.PRNGKey(1), data)
```
I've found this solution while comparing documentation code to "Code 4.53" by @fehiepsi from [here](https://github.com/fehiepsi/rethinking-numpyro/blob/master/notebooks/04_geocentric_models.ipynb ) 

---

My full model code:
```python
import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)

import pandas as pd

import jax.numpy as jnp
from jax import random

import numpyro
import numpyro.distributions as dist
import numpyro.optim as optim
from numpyro.infer import SVI, Trace_ELBO, Predictive
from numpyro.infer.autoguide import AutoNormal

# INFO: load data for problems from the official repo
data_uri = "https://raw.githubusercontent.com/rmcelreath/rethinking/master/data/Howell1.csv"
df_dev = pd.read_csv(data_uri, sep=";")
df_dev = df_dev[df_dev["age"] >= 18]

def model(df_dev):
    alpha = numpyro.sample("alpha", dist.Normal(60, 10))
    beta = numpyro.sample("beta", dist.LogNormal(0, 1))
    sigma = numpyro.sample("sigma", dist.Uniform(0, 10))
    mu = numpyro.deterministic("mu", alpha + beta * (df_dev["height"] - df_dev["height"].mean()).values)
    numpyro.sample("W", dist.Normal(mu, sigma), obs=df_dev["weight"].values)



# quadratic approximation part
guide = AutoNormal(model)
svi = SVI(model, guide, optim.Adam(1), Trace_ELBO(), df_dev=df_dev)
svi_result = svi.run(random.PRNGKey(0), 1000)

params = svi_result.params

# display summary of quadratic approximation
samples = guide.sample_posterior(random.PRNGKey(1), params, (1000,))
numpyro.diagnostics.print_summary(samples, prob=0.89, group_by_chain=False)

# INFO: bad
predictive = Predictive(guide, params=params, posterior_samples=samples)
samples_predictive = predictive(random.PRNGKey(1), df_dev)
print(samples_predictive.keys())

# INFO: good
predictive = Predictive(guide.model, params=params, posterior_samples=samples)
samples_predictive = predictive(random.PRNGKey(1), df_dev)
print(samples_predictive.keys())

```
---
Please tell me what do you think


Discussion:
Typically, to get posterior samples, you can use Predictive(guide, ...). To get posterior predictive, you can use Predictive(model, guide=guide, ...). Predictive(model, ...) will give you prior predictive. I think what you want is posterior predictive so please use Predictive(model, guide=guide,... It would be nice to improve the docs to get posterior predictive. Do you want to contribute? :)
@fehiepsi correct, i was after posterior predictive. And yes i could work on PR

I think what you said should be mentioned in the `svi.py` and `util.py`.
Particularly `util.py` suggests the `Predictive(model, posterior_samples=posterior_samples, ...)` recipe (granted because MCMC doesn't imply any `guide`)  
https://github.com/pyro-ppl/numpyro/blob/5d9e033297532ae1a81a3e762ce7250aed5ab269/numpyro/infer/util.py#L839-L842



And the `svi,py` example, to me at least, is quite low-level and cryptic and it provides more questions than answers 🙂 I'd personally redo it into single simple/minimalistic example (e.g. w/o custom guide or plate) or simple-to-advanced series of examples (where custom guide & plate get introduced), like they do in `Pandas` docs. Speaking from the fact that i don't even understand the meaning of neither custom guide nor plate part 
https://github.com/pyro-ppl/numpyro/blob/a35e9254bfbf2a1a84138670f48bdc45c94f829e/numpyro/infer/svi.py#L113-L115

---
For a quick fix i'd change the `svi.py` example above into something like:

```python
>>> # get posterior samples
>>> predictive = Predictive(model=guide, params=params, num_samples=1000)
>>> samples = predictive(random.PRNGKey(1), data)

>>> # get posterior predictive samples, approach 1: from model
>>> predictive = Predictive(model=model, guide=guide, params=params, num_samples=1000)
>>> samples = predictive(random.PRNGKey(1), data)

>>> # get posterior predictive samples, approach 2: from samples
>>> posterior_samples = guide.sample_posterior(random.PRNGKey(1), params, (1000,))
>>> predictive = Predictive(model=model, guide=guide, params=params, posterior_samples=posterior_samples)
>>> samples = predictive(random.PRNGKey(1), data)

>>> # get prior predictive samples
>>> predictive = Predictive(model=model, params=params, num_samples=1000)
>>> samples = predictive(random.PRNGKey(1), data)

```
(Please correct me becase there's something excessive, I sense it 🧛‍♂️)

Does it make sense?

Yes, that makes the svi docs much accessible. Just two suggestions - Maybe move the posterior_samples to the first section and comment it as an alternative way to draw posterior samples from auto guides. We should rename `samples` to `posterior_samples` to avoid confusion. - To get posterior/prior predictive, we should set data=None in predictive(..., data=None)
Can't run "prior predictive" without data/with data=None, it errors, because the model demands its parameters.

Also I wonder why `"W"`, when it's present in `samples`, is always equal to N identical rows of `data`. Like it's not sampled from the `dist.Normal(mu, sigma)`, it's just outputting a matrix of `df_dev["weight"].values` copies. Is this intended?
Also 

```python
>>> # get posterior samples
>>> predictive = Predictive(model=guide, params=params, num_samples=1000)
>>> samples = predictive(random.PRNGKey(1))
```
is not equivalent to  
```python
samples = guide.sample_posterior(random.PRNGKey(1), params, (1000,))
```

I have a case where 1 outputs constant values and 2 outputs proper samples. Am I missing something?
`# get posterior predictive samples, approach 1: from model` is quite prone to this aswell in my expereience (contrary to `# get posterior predictive samples, approach 2: from samples`)

> Can't run "prior predictive" without data/with data=None

I think in the doc, you can use `with numpyro.plate("N", 10):` rather than `data.shape[0]` (so that `data` can be None, i.e. unobserved).

> why "W", when it's present in samples, is always equal to N identical rows of data

You need to set `obs=None` in the model. One way to do it is
```
def model(df_dev, observed=True):
    ...
    numpyro.sample("W", dist.Normal(mu, sigma), obs=df_dev["weight"].values if observed else None)

...
samples = predictive(random.PRNGKey(1), df_dev, observed=False)
```

> is not equivalent to

They should be similar in functionality. Using `guide.sample_posterior` might return for you `deterministic` sites. You can achieve the sample purpose with `Predictive(..., return_sites=list_of_desired_sites)`. I'm not sure why you got **constant** values. Could you provide some code?

> I think in the doc, you can use with numpyro.plate("N", 10): rather than data.shape[0] (so that data can be None, i.e. unobserved).

I think introducing a plate abstraction to a simple example of basics does more harm than good. And sampling prior is a basic functionality of a bayesian model.

So i've came up with this:
```python
# get pripr predictive samples
# set `params=None` explicitly to express that they're not used, since it's prior predictive 
predictive = Predictive(model=model, params=None, num_samples=1000)
# no-op data `df_dev * 0` to show explicitly that it's only for compatibility and does not affect inference
# optional `observed=False` to make the target variable sampled from prior aswell (otherwise it'll output the no-op data)
samples = predictive(random.PRNGKey(1), df_dev * 0, observed=False)
samples
```

Full example:
```python
# setup
import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)

import pandas as pd

import jax.numpy as jnp
from jax import random

import numpyro
import numpyro.distributions as dist
import numpyro.optim as optim
from numpyro.infer import SVI, Trace_ELBO, Predictive
from numpyro.infer.autoguide import AutoNormal

# INFO: load data for problems from the official repo
data_uri = "https://raw.githubusercontent.com/rmcelreath/rethinking/master/data/Howell1.csv"
df_dev = pd.read_csv(data_uri, sep=";")
df_dev = df_dev[df_dev["age"] >= 18]

def model(df_dev, observed=True):
    alpha = numpyro.sample("alpha", dist.Normal(60, 10))
    beta = numpyro.sample("beta", dist.LogNormal(0, 1))
    sigma = numpyro.sample("sigma", dist.Uniform(0, 10))
    mu = numpyro.deterministic("mu", alpha + beta * (df_dev["height"] - df_dev["height"].mean()).values)
    numpyro.sample("W", dist.Normal(mu, sigma), obs=df_dev["weight"].values if observed else None)


# quadratic approximation part
guide = AutoNormal(model)
svi = SVI(model, guide, optim.Adam(1), Trace_ELBO(), df_dev=df_dev)
svi_result = svi.run(random.PRNGKey(0), 1000)

params = svi_result.params

# get pripr predictive samples
# set `params=None` explicitly to express that they're not used, since it's prior predictive 
predictive = Predictive(model=model, params=None, num_samples=1000)
# no-op data to show explicitly that it's only for compatibility and does not affect inference
# optional `observed=False` to make the target variable sampled from prior aswell (otherwise it'll output the no-op data)
samples = predictive(random.PRNGKey(1), df_dev * 0, observed=False)
samples
```

---

> You need to set obs=None in the model. One way to do it is

That trick is very neat, thanks!

---

> They should be similar in functionality. Using guide.sample_posterior might return for you deterministic sites. You can achieve the sample purpose with Predictive(..., return_sites=list_of_desired_sites). I'm not sure why you got constant values. Could you provide some code?

```python

# INFO: snippet for GH#1309
# https://github.com/pyro-ppl/numpyro/issues/1309#issuecomment-1025235616

from pprint import pprint

import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)

import pandas as pd

import jax.numpy as jnp
from jax import random

import numpyro
import numpyro.distributions as dist
import numpyro.optim as optim
from numpyro.infer import SVI, Trace_ELBO, Predictive
from numpyro.infer.autoguide import AutoNormal

# INFO: load data for problems from the official repo
data_uri = "https://raw.githubusercontent.com/rmcelreath/rethinking/master/data/Howell1.csv"
df_dev = pd.read_csv(data_uri, sep=";")
df_dev = df_dev[df_dev["age"] >= 18]


# INFO: simplest model, now with data, r4.16, fig4.4
def model(df_dev):
    mu = numpyro.sample("mu", dist.Normal(178, 20))
    sigma = numpyro.sample("sigma", dist.Uniform(0, 50))

    # INFO: note: we supply data here
    h = numpyro.sample("h", dist.Normal(mu, sigma), obs=df_dev["height"].values)


guide = AutoLaplaceApproximation(model)
svi = SVI(model=model, guide=guide, optim=optim.Adam(0.01), loss=Trace_ELBO(), df_dev=df_dev)
num_steps = 1000
svi_result = svi.run(random.PRNGKey(0), num_steps)

params = svi_result.params

# display summary of quadratic approximation
# INFO: `sample_shape` arg defines the size of the sample, not svi sample size or dataset size ()
guide_samples = guide.sample_posterior(random.PRNGKey(1), params, (10_000,))
print("\nguide.sample_posterior")
pprint(guide_samples)

num_samples = 10_000  # will produce a matrix of identical rows
# num_samples = 1
# INFO: sample posterior - constant values wow much useful
predictive_sampler = Predictive(model=guide, params=params, num_samples=num_samples)
predictive_sampler_samples = predictive_sampler(random.PRNGKey(3))
print("\nPredictive.__call__")
pprint(predictive_sampler_samples)


def test_identical_values(samples, variable):
    try:
        a_full = jnp.full_like(samples[variable], fill_value=samples[variable][0])
        assert jnp.array_equal(samples[variable], a_full)
    except:
        print("array values are not identical")
    else:
        print("array values are identical!")


test_identical_values(predictive_sampler_samples, "sigma")
# expected: array values are not identical
# actual: array values are identical!
test_identical_values(guide_samples, "sigma")
# expected: array values are not identical
# actual: array values are not identical
```


Ok i've redesigned the whole "different samplings from SVI" example because my previous one looked stupid and had redundant arguments. And updated with new tricks. I use this one as a cheat sheet for myself.

Has all basic kinds of sampling and samples `W` variable directly (so no need to reconstruct the whole formula/distribution just to plot the damn thing)

```python

# INFO: load data for problems from the official repo
data_uri = "https://raw.githubusercontent.com/rmcelreath/rethinking/master/data/Howell1.csv"
df_dev = pd.read_csv(data_uri, sep=";")

df_train = df_dev[df_dev["age"] >= 18]
X_train = df_dev.loc[:, ["height"]]
y_train = df_dev.loc[:, ["weight"]]

df_test = pd.DataFrame(data={"height": [140, 160, 175], "weight": [None, None, None]})
X_test = df_test.loc[:, ["height"]]
y_test = df_test.loc[:, ["weight"]]


def model(X: pd.DataFrame, y: pd.DataFrame, observed=True):
    alpha = numpyro.sample("alpha", dist.Normal(60, 10))
    beta = numpyro.sample("beta", dist.LogNormal(0, 1))
    sigma = numpyro.sample("sigma", dist.Uniform(0, 10))
    # INFO: subtract X_train explicitly to secure non-leaky inference
    mu = numpyro.deterministic("mu", alpha + beta * (X["height"] - X_train["height"].mean()).values)
    numpyro.sample("W", dist.Normal(mu, sigma), obs=y["weight"].values if observed else None)


# quadratic approximation part
guide = AutoNormal(model)
svi = SVI(model, guide, optim.Adam(1), Trace_ELBO(), X=X_train, y=y_train)
svi_result = svi.run(random.PRNGKey(0), 1000)

params = svi_result.params

# display summary of quadratic approximation
samples = guide.sample_posterior(random.PRNGKey(1), params, (1000,))
numpyro.diagnostics.print_summary(samples, prob=0.89, group_by_chain=False)

# ---

# get posterior samples, approach 1: from distribution
num_samples = 1000
dist_posterior = Predictive(model=guide, params=params, num_samples=num_samples)
samples_posterior = dist_posterior(random.PRNGKey(1))
samples_posterior

# get posterior samples, approach 2: from guide
num_samples = 1000
samples_posterior = guide.sample_posterior(random.PRNGKey(1), params=params, sample_shape=(num_samples,))
samples_posterior

# ---

# get posterior predictive samples, approach 1: from samples
num_samples = 1000
samples_posterior = guide.sample_posterior(random.PRNGKey(1), params=params, sample_shape=(num_samples,))
dist_posterior_predictive = Predictive(model=model, guide=guide, params=params, posterior_samples=samples_posterior)
samples_posterior_predictive = dist_posterior_predictive(random.PRNGKey(1), X=X_test, y=None, observed=False)
samples_posterior_predictive

# get posterior predictive samples, approach 2: from model
# approach 2, for test data
num_samples = 1000
dist_posterior_predictive = Predictive(model=model, guide=guide, params=params, num_samples=num_samples)
samples_posterior_predictive = dist_posterior_predictive(random.PRNGKey(1), X=X_test, y=None, observed=False)
samples_posterior_predictive

# approach 2, for train data (e.g. to asess model fit)
num_samples = 1000
dist_posterior_predictive = Predictive(model=model, guide=guide, params=params, num_samples=num_samples)
samples_posterior_predictive = dist_posterior_predictive(random.PRNGKey(1), X=X_train, y=None, observed=False)
samples_posterior_predictive

# approach 2, for counterfactual data (aka simulation, generative modeling)
X_gen = pd.DataFrame({"height": np.arange(10, 200, 10)})

num_samples = 1000
dist_posterior_predictive = Predictive(model=model, guide=guide, params=params, num_samples=num_samples)
samples_posterior_predictive = dist_posterior_predictive(random.PRNGKey(1), X=X_gen, y=None, observed=False)
samples_posterior_predictive

# ---

# get prior predictive samples
# for test data
num_samples = 1000
dist_prior_predictive = Predictive(model=model, params=None, num_samples=num_samples)
samples_prior_predictive = dist_prior_predictive(random.PRNGKey(1), X=X_test, y=None, observed=False)
samples_prior_predictive

```

The `X["height"] - X_train["height"].mean()` part is ugly but i didn't want to bring in `sklearn` library to a simple example fight, or create additional args, or perform awkward transformations beforehand. In a real project I'd likely just use `sklearn` anyway.
Re constant values: AutoLaplaceApproximation uses a delta guide (which does MAP estimate, see [docs](https://num.pyro.ai/en/stable/autoguide.html#numpyro.infer.autoguide.AutoLaplaceApproximation)) so if you draw samples from it, you will get constant values.
Re plate: yeah, it might be confusing at first if we don't know about [plate diagram](https://en.wikipedia.org/wiki/Plate_notation)  - it will be pretty handy for complicated modeling, where you can move variables around just like you move circles inside/outside rectangle boxes - no need to specify the shape correctly. Sometimes, it will be easy to make mistakes without plates - let's take the example in SVI docs, without plate, predictive with 1000 samples will do the wrong thing
```
def model(data):
    f = numpyro.sample("latent_fairness", dist.Beta(10, 10))
    numpyro.sample("obs", dist.Bernoulli(f), obs=data)
```
because when `data=None`, `obs` will have the shape `(1000,)` rather than the desired one `(1000, 10)`. Other PPLs resort `plate` to something like
```
numpyro.sample("obs", dist.Bernoulli(f).expand([10]), obs=data)
```
or
```
numpyro.sample("obs", dist.Bernoulli(f), obs=data, sample_shape=(10,))
```
They might work in NumPyro but personally, I found that using `plate` is handier.

Re your new model: It looks pretty nice. Just a very minor comment - maybe you can simplify things a bit by not using data frames for `y` (so no need for `observed` flag)
```
def model(X, y=None):
    ...
    return numpyro.sample("obs", d, obs=y)
```
> Re constant values: AutoLaplaceApproximation uses a delta guide (which does MAP estimate, see [docs](https://num.pyro.ai/en/stable/autoguide.html#numpyro.infer.autoguide.AutoLaplaceApproximation)) so if you draw samples from it, you will get constant values.

Hm, even tho I had read the doc, it was not obvious to me that these are explicitly point estimates. Until I found how does Delta distribution look like 😵

Re plates
I'm a tad familiar with plate notation. Now, thinking of it again, from a vectorized code standpoint, it might indeed be a handy thing - like a shell for data which is empty when we don't supply it. I have to think about it.
Btw, can I ask you why you(core devs) chose that specific example for SVI api (the 'latent_fairness' one)? 

Re model
Thanks, I like pandas for workflow, but ofc the library code is good to be JAX-centric. I'll tailor the examples for PR
> it was not obvious to me that these are explicitly point estimates

In your PR, could you update docs of LaplaceApproximation to: `... construct a MAP guide (i.e. point estimates to maximize the posterior)`

> chose that specific example for SVI api

I can't recall. It seems to just be a random choice - we want an example that has `param` `sample` `plate` statements. :D 
@ColdTeapot273K Do you still want to take a stab at this? Your suggestions would be nice improvements for the docs.
@fehiepsi hey thanks for the heads up. 

I shall try to submit a PR this week✊

Btw wrote a higher level micro framework to handle this boilerplate in a sklearn fashion, leveraging xarray lib among other things (kinda like arviz but minimalist). Ill show it later maybe
"So that was a f* lie"

Sorry, had to attend emergency matters popping up due to the recent global events🥵

I aim to contrib this weekend

Resolving patch:
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
--- a/numpyro/infer/autoguide.py
+++ b/numpyro/infer/autoguide.py
@@ -1647,8 +1647,8 @@ class AutoLaplaceApproximation(AutoContinuous):
     Laplace approximation (quadratic approximation) approximates the posterior
     :math:`\log p(z | x)` by a multivariate normal distribution in the
     unconstrained space. Under the hood, it uses Delta distributions to
-    construct a MAP guide over the entire (unconstrained) latent space. Its
-    covariance is given by the inverse of the hessian of :math:`-\log p(x, z)`
+    construct a MAP (i.e. point estimate) guide over the entire (unconstrained) latent
+    space. Its covariance is given by the inverse of the hessian of :math:`-\log p(x, z)`
     at the MAP point of `z`.
 
     Usage::
diff --git a/numpyro/infer/svi.py b/numpyro/infer/svi.py
--- a/numpyro/infer/svi.py
+++ b/numpyro/infer/svi.py
@@ -95,7 +95,7 @@ class SVI(object):
 
         >>> def model(data):
         ...     f = numpyro.sample("latent_fairness", dist.Beta(10, 10))
-        ...     with numpyro.plate("N", data.shape[0]):
+        ...     with numpyro.plate("N", data.shape[0] if data is not None else 10):
         ...         numpyro.sample("obs", dist.Bernoulli(f), obs=data)
 
         >>> def guide(data):
@@ -110,9 +110,15 @@ class SVI(object):
         >>> svi_result = svi.run(random.PRNGKey(0), 2000, data)
         >>> params = svi_result.params
         >>> inferred_mean = params["alpha_q"] / (params["alpha_q"] + params["beta_q"])
+        >>> # use guide to make predictive
+        >>> predictive = Predictive(model, guide=guide, params=params, num_samples=1000)
+        >>> samples = predictive(random.PRNGKey(1), data=None)
         >>> # get posterior samples
         >>> predictive = Predictive(guide, params=params, num_samples=1000)
-        >>> samples = predictive(random.PRNGKey(1), data)
+        >>> posterior_samples = predictive(random.PRNGKey(1), data=None)
+        >>> # use posterior samples to make predictive
+        >>> predictive = Predictive(model, posterior_samples, params=params, num_samples=1000)
+        >>> samples = predictive(random.PRNGKey(1), data=None)
 
     :param model: Python callable with Pyro primitives for the model.
     :param guide: Python callable with Pyro primitives for the guide


Test patch: