Allow pickle autoguide

[fehiepsi]

Resolves #1160

TODO:

• allow pickling unpack_latent
• allow pickling postprocess_fn
• add tests

[fehiepsi]

This follows Pyro patch for PyTorch transforms.

[fehiepsi]

@lumip This supposes to fix your pickling issue. Could you double-check? Thanks!

[fehiepsi]

plates is generated during execution, so we don't need to cache. Actually, we typically call the guide under jit, so plates will hold abstract values which cannot be pickled.

[fritzo]

Test looks good 👍 It would be nice to be able to unpickle and use a guide without the model (via guide()), but I guess this is a separate issue.

[fehiepsi]

Thanks for reviewing, @fritzo! You can use guide to sample posteriors by using either

guide.sample_posterior(rng_key, svi_result.params, sample_shape=(num_samples,))

or

Predictive(guide, params=svi_result.params, num_samples=num_samples)(rng_key)

We need it because we need params, rng_key to be able to call guide().

Using Predictive(model, guide, params) is unnecessary. It is just a shortcut for

posterior_samples = guide.sample_posterior(...)
x = Predictive(model, posterior_samples, ...)(rng_key, ...)['x']

cc @lumip

[lumip]

Maybe better to assert that pickled_guide returns the same samples as guide instead of just verifying that all sites are present?

[lumip]

This seems to work well in some simple tests I ran. Unfortunately I can't right now verify that in our more elaborate settings but I may get around to that tomorrow.

Unfortunately I'm not all too familiar with the code that was changed, so I cannot review those changes in detail.

Edit: Have now tried it in a more complicated environment and there I ran into the issue that we have some additional decorators around our model which prevent me from pickling the guard (I get a AttributeError: Can't pickle local object 'guard_model.<locals>.wrapped_model' where the guard_model is our decorator). So I guess the pickling approach works, but only if the user provides a model that is pickle-able as well.

[fehiepsi]

Thanks for reviewing, @lumip! I think we can just remove model from a pickled guide (if the guide already has prototype trace). Let me address that. For now, you can just set guide.model = None or delete that attribute before pickling.

[fehiepsi]

@lumip Have been thinking about this for a while, it seems that it's better to let users decide to pickle guide with model or guide without model. If we remove model automatically, someone might be surprised that they can't access model attribute from the guide anymore. :(