Repository: pyro-ppl/numpyro

Issue:
AutoGuides do not permit sampling using only learned parameters
## The Goal
Our setting requires that we separate sampling from a learned variational posterior some time after and in a separate script/program than the inference and *without having access to the training data* any more.

## The Problem
We use NumPyro's SVI algorithm with the `AutoDiagonalNormal` guide to infer a variational approximation of our model posterior. We then store the learned parameters of the guide.

At a later time, to sample from the learned posterior, we load the stored parameters, set up our model and the `AutoDiagonalNormal` guide and then call `Predictive` to sample from the data. However, rather than giving us samples from our approximated posterior, this results in one of several possible errors (depending on the internals of the model).

Example code:
```Python
import numpy as np
import jax
import numpyro
import numpyro.distributions as dists
from numpyro.infer import SVI, Predictive
from numpyro.infer.autoguide import AutoDiagonalNormal
from numpyro.primitives import sample, plate

d = 2
np.random.seed(12345)

x = np.random.poisson(1.0, size=(100,))

def model(x, N):
    rate = sample('param', dists.Gamma(1., 1.))
    batch_size = len(x) if x is not None else None
    with plate('batch', N, batch_size):
        sample('x', dists.Poisson(rate), obs=x)

guide = AutoDiagonalNormal(model)

# run inference
optim = numpyro.optim.Adam(1e-2)
svi = SVI(model, guide, optim, numpyro.infer.Trace_ELBO())
state = svi.init(jax.random.PRNGKey(0), x, len(x))
state = jax.lax.fori_loop(0, 10000, lambda i, state: svi.update(state, x, len(x))[0], state)

# get params
params = svi.get_params(state)

## sampling directly after inference works fine
samples = Predictive(model, guide=guide, params=params, num_samples=1, return_sites=['param', 'x'])(
    jax.random.PRNGKey(1), None, 1
)
print(samples)

## sampling with fresh guide does not work
guide = AutoDiagonalNormal(model) # fresh setup of guide to emulate script without inference

samples = Predictive(model, guide=guide, params=params, num_samples=1, return_sites=['param', 'x'])(
    jax.random.PRNGKey(1), None, 1
)
print(samples)
```
The second call to `Predictive` raises a `RuntimeError: MCMC only supports continuous sites or discrete sites with enumerate support, but got ExpandedDistribution.` (but really it has issues with the `Poisson`). Using other distributions will results in different errors.

## A bit of digging into this
We did a bit of digging and have come to the conclusion that this is caused by how the autoguides work: When they are first used, they will trace through the model to determine how to convert samples from an unconstrained domain to whatever is required in the model. It does this for all sample sites that are not observations. This is problematic in our case since we only want to sample from the learned posterior and therefore do not at any point fix observations in the model. The autoguide therefore treats the final sample sites corresponding to observations as something it needs to optimise (even though it will never have to do that).

## Ways forward?
We believe this is quite an oversight: Since the learned parameters fully describe the variational posterior mathematically, saving and reloading those should suffice to sample from the posterior in the API as well, without having to pass data into the model again (which is also a promise in the docstring of `AutoDiagonalNormal`': "The guide does not depend on the model's ``*args, **kwargs``.")

A workaround for us currently is to simply pass in some dummy data to initialize the guide, then go forward to sample from it using the learned parameters. However, this is somewhat cumbersome and unintuitive.

We think that ideally `Predictive` or the autoguide should be made smart enough to detect this case (maybe with an additional input for autoguide on which sites should be considered observations even when no data is given) and handle it.

At least it should be possible to store and load the state of the autoguide alongside the parameters. We tried to see whether this is possible at the moment. It didn't seem to be the case, but maybe we missed something.

The error messages for this should also be improved, as they don't really point to the actual underlying problem and are not consistent (see below for more examples).

What do you think? Or is there already a way to achieve our purpose using NumPyro that we have not thought of?

## Errors for Other Distributions
### Categorical distribution:
```
x = np.random.multinomial(1, [.5, .25, .125, .125], size=(100,)) * np.array([1,2,3,4])
x = np.sum(x, axis=-1)

# Categorical
def model(x, N):
    ps = sample('param', dists.Dirichlet(np.ones(4,)))
    batch_size = len(x) if x is not None else None
    with plate('batch', N, batch_size):
        sample('x', dists.Categorical(ps), obs=x)
```
This raises a `ImportError: Looking like you want to do inference for models with discrete latent variables. This is an experimental feature. You need to install funsor to be able to use this feature. It can be installed with pip install funsor.`

### Normal distribution:
```
x = np.random.normal(10, 3, size=(100, 2))
x[:,-1] *= -1

# Normal
def model(x, N):
    loc = sample('param', dists.Normal(0, 100), sample_shape=(d,))
    scale = np.array([2., 2.])

    batch_size = len(x) if x is not None else None
    with plate('batch', N, batch_size):
        sample('x', dists.Normal(loc, scale).to_event(1), obs=x)
```
This raises a `TypeError: reshape total size must be unchanged, got new_sizes (1, 2) for shape (0,).`.


Discussion:
In Pyro you can just call `guide()` to draw samples of latent variables: there is no need for `Predictive` except to perform deterministic computations in the model. Does that work for you?  EDIT of course you'll need to store the guide between scripts, but when you load the guide you no longer need access to the model or data.
Thanks @fritzo for pointing that out. However, the errors currently occur directly from the invocation of `guide` that happens internally in `Predictive`, so avoiding `Predictive` does unfortunately not solve the issue.

Could you specify what you mean by "you'll need to store the guide"? Are you suggesting to pickle the entire class instance? Our initial thought was that the posterior is fully described by the parameters learned for the guide and the implementation of the model. As explained above, this is not the case. We have not found a way of extracting the state of the autoguide for storage either, so I guess the entire question is about how to store an accurate representation of the guide (in a way that doesn't break with the next Python/NumPyro/JAX update).
@fehiepsi how do you serialize guides in NumPyro?  It would be nice if we could get pickle working, say via `__getstate__`/`__setstate__`.  In Pyro we simply:
```py
torch.save(guide, filename)
...
guide = torch.load(filename)
```
I'll look into how to store state (though I'm not sure which state we need to store). @lumip's way to setup the guide with dummy data seems to be a good solution. It lets the guide know a bit information of the model (without that, model is just a function - we don't have a way to trace/interpret how the model works from model's source code, so even when we provide `params`, the guide won't know how to use that `params` variable.
I guess current solution is
```
guide = AutoDiagonalNormal(model)
with numpyro.handlers.seed(rng_seed=0):
    guide._setup_prototype(x, len(x))
    # or just guide(x, len(x))
```
And we want to support the following pattern
```
pickle.dump(guide, filename)
guide = pickle.load(filename)
```
so that @lumip can use the guide in another script without passing model's args, kwargs. It seems to be non-trivial at my first glance.

Resolving patch:
diff --git a/numpyro/distributions/transforms.py b/numpyro/distributions/transforms.py
--- a/numpyro/distributions/transforms.py
+++ b/numpyro/distributions/transforms.py
@@ -106,6 +106,16 @@ def inverse_shape(self, shape):
         """
         return shape
 
+    # Allow for pickle serialization of transforms.
+    def __getstate__(self):
+        attrs = {}
+        for k, v in self.__dict__.items():
+            if isinstance(v, weakref.ref):
+                attrs[k] = None
+            else:
+                attrs[k] = v
+        return attrs
+
 
 class _InverseTransform(Transform):
     def __init__(self, transform):
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
--- a/numpyro/infer/autoguide.py
+++ b/numpyro/infer/autoguide.py
@@ -4,6 +4,7 @@
 # Adapted from pyro.infer.autoguide
 from abc import ABC, abstractmethod
 from contextlib import ExitStack
+from functools import partial
 import warnings
 
 import numpy as np
@@ -11,7 +12,6 @@
 import jax
 from jax import grad, hessian, lax, random, tree_map
 from jax.experimental import stax
-from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 
 import numpyro
@@ -104,6 +104,11 @@ def _create_plates(self, *args, **kwargs):
                 )
         return self.plates
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state.pop("plates", None)
+        return state
+
     @abstractmethod
     def __call__(self, *args, **kwargs):
         """
@@ -450,6 +455,34 @@ def median(self, params):
         return locs
 
 
+def _unravel_dict(x_flat, shape_dict):
+    """Return `x` from the flatten version `x_flat`. Shape information
+    of each item in `x` is defined in `shape_dict`.
+    """
+    assert jnp.ndim(x_flat) == 1
+    assert isinstance(shape_dict, dict)
+    x = {}
+    curr_pos = next_pos = 0
+    for name, shape in shape_dict.items():
+        next_pos = curr_pos + int(np.prod(shape))
+        x[name] = x_flat[curr_pos:next_pos].reshape(shape)
+        curr_pos = next_pos
+    assert next_pos == x_flat.shape[0]
+    return x
+
+
+def _ravel_dict(x):
+    """Return the flatten version of `x` and shapes of each item in `x`."""
+    assert isinstance(x, dict)
+    shape_dict = {}
+    x_flat = []
+    for name, value in x.items():
+        shape_dict[name] = jnp.shape(value)
+        x_flat.append(value.reshape(-1))
+    x_flat = jnp.concatenate(x_flat) if x_flat else jnp.zeros((0,))
+    return x_flat, shape_dict
+
+
 class AutoContinuous(AutoGuide):
     """
     Base class for implementations of continuous-valued Automatic
@@ -474,7 +507,8 @@ class AutoContinuous(AutoGuide):
 
     def _setup_prototype(self, *args, **kwargs):
         super()._setup_prototype(*args, **kwargs)
-        self._init_latent, unpack_latent = ravel_pytree(self._init_locs)
+        self._init_latent, shape_dict = _ravel_dict(self._init_locs)
+        unpack_latent = partial(_unravel_dict, shape_dict=shape_dict)
         # this is to match the behavior of Pyro, where we can apply
         # unpack_latent for a batch of samples
         self._unpack_latent = UnpackTransform(unpack_latent)


Test patch:
diff --git a/test/test_pickle.py b/test/test_pickle.py
--- a/test/test_pickle.py
+++ b/test/test_pickle.py
@@ -3,6 +3,7 @@
 
 import pickle
 
+import numpy as np
 import pytest
 
 from jax import random, test_util
@@ -16,10 +17,13 @@
     MCMC,
     NUTS,
     SA,
+    SVI,
     BarkerMH,
     DiscreteHMCGibbs,
     MixedHMC,
+    Predictive,
 )
+from numpyro.infer.autoguide import AutoDelta, AutoDiagonalNormal, AutoNormal
 
 
 def normal_model():
@@ -59,3 +63,31 @@ def test_pickle_hmcecs():
     mcmc.run(random.PRNGKey(0))
     pickled_mcmc = pickle.loads(pickle.dumps(mcmc))
     test_util.check_close(mcmc.get_samples(), pickled_mcmc.get_samples())
+
+
+def poisson_regression(x, N):
+    rate = numpyro.sample("param", dist.Gamma(1.0, 1.0))
+    batch_size = len(x) if x is not None else None
+    with numpyro.plate("batch", N, batch_size):
+        numpyro.sample("x", dist.Poisson(rate), obs=x)
+
+
+@pytest.mark.parametrize("guide_class", [AutoDelta, AutoDiagonalNormal, AutoNormal])
+def test_pickle_autoguide(guide_class):
+    x = np.random.poisson(1.0, size=(100,))
+
+    guide = guide_class(poisson_regression)
+    optim = numpyro.optim.Adam(1e-2)
+    svi = SVI(poisson_regression, guide, optim, numpyro.infer.Trace_ELBO())
+    svi_result = svi.run(random.PRNGKey(1), 3, x, len(x))
+    pickled_guide = pickle.loads(pickle.dumps(guide))
+
+    predictive = Predictive(
+        poisson_regression,
+        guide=pickled_guide,
+        params=svi_result.params,
+        num_samples=1,
+        return_sites=["param", "x"],
+    )
+    samples = predictive(random.PRNGKey(1), None, 1)
+    assert set(samples.keys()) == {"param", "x"}