Repository: pyro-ppl/pyro

Issue:
Make dynamic distribution mixin wrappers in torch.py picklable
We load torch distributions and create distribution mixin classes dynamically in `torch.py`. These distribution instances are not picklable however.

```python
>>> pickle.dumps(Pareto(1., 1.))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
_pickle.PicklingError: Can't pickle <class 'pyro.distributions.torch._PyroDist'>: it's not the same object as pyro.distributions.torch._PyroDist
```
We may need to provide more help to `pickle` to be able pickle / unpickle these instances. Note that this is not currently an issue with most distribution classes since they have a small wrapper in Pyro (containing `.expand` method), but will be in `pyro==0.3` where we are constructing these wrappers automatically.

### Why do we need this?

This will be useful in the multiprocessing context where workers send back traces to the main process. All elements in the trace (inc. distribution instances) should be picklable in that case. An alternative design would be to not send back the full trace, but only the empirical distribution itself, but this is worth exploring only if the overhead of trace serialization/deserialization is considerable.


Discussion:
I'd guess we can simply set `type(...).__name__` to e.g. `Pareto` and set that as a module level variable in the `pyro.distributions.torch` module.
@fritzo - This is only an issue if we have distribution classes in PyTorch that we would like to dynamically import. 

If `Pareto` is available as a module level variable in `pyro.distributions.torch`, then pickle wouldn't have an issue. Declaring all distribution classes in the `torch` module upfront  is definitely one solution (and maybe better than being too clever). But I might have misunderstood your solution - would it work with dynamic importing?

Resolving patch:
diff --git a/pyro/distributions/torch.py b/pyro/distributions/torch.py
--- a/pyro/distributions/torch.py
+++ b/pyro/distributions/torch.py
@@ -313,11 +313,8 @@ def expand(self, batch_shape):
     try:
         _PyroDist = locals()[_name]
     except KeyError:
-
-        class _PyroDist(_Dist, TorchDistributionMixin):
-            pass
-
-        _PyroDist.__name__ = _name
+        _PyroDist = type(_name, (_Dist, TorchDistributionMixin), {})
+        _PyroDist.__module__ = __name__
         locals()[_name] = _PyroDist
 
     _PyroDist.__doc__ = '''


Test patch:
diff --git a/tests/distributions/test_pickle.py b/tests/distributions/test_pickle.py
new file mode 100644
--- /dev/null
+++ b/tests/distributions/test_pickle.py
@@ -0,0 +1,80 @@
+from __future__ import absolute_import, division, print_function
+
+import inspect
+
+import pytest
+import six.moves.cPickle as pickle
+import torch
+
+import pyro.distributions as dist
+from pyro.distributions.torch_distribution import TorchDistributionMixin
+from tests.common import xfail_param
+
+# Collect distributions.
+BLACKLIST = [
+    dist.TorchDistribution,
+    dist.ExponentialFamily,
+    dist.OMTMultivariateNormal,
+]
+XFAIL = {
+    dist.Gumbel: xfail_param(dist.Gumbel, reason='cannot pickle weakref'),
+}
+DISTRIBUTIONS = [d for d in dist.__dict__.values()
+                 if isinstance(d, type)
+                 if issubclass(d, TorchDistributionMixin)
+                 if d not in BLACKLIST]
+DISTRIBUTIONS.sort(key=lambda d: d.__name__)
+DISTRIBUTIONS = [XFAIL.get(d, d) for d in DISTRIBUTIONS]
+
+# Provide default args if Dist(1, 1, ..., 1) is known to fail.
+ARGS = {
+    dist.AVFMultivariateNormal: [torch.zeros(3), torch.eye(3), torch.rand(2, 4, 3)],
+    dist.Bernoulli: [0.5],
+    dist.Binomial: [2, 0.5],
+    dist.Categorical: [torch.ones(2)],
+    dist.Delta: [torch.tensor(0.)],
+    dist.Dirichlet: [torch.ones(2)],
+    dist.GaussianScaleMixture: [torch.ones(2), torch.ones(3), torch.ones(3)],
+    dist.Geometric: [0.5],
+    dist.Independent: [dist.Normal(torch.zeros(2), torch.ones(2)), 1],
+    dist.LowRankMultivariateNormal: [torch.zeros(2), torch.ones(2, 2), torch.ones(2)],
+    dist.MaskedMixture: [torch.tensor([1, 0]).byte(), dist.Normal(0, 1), dist.Normal(0, 2)],
+    dist.MixtureOfDiagNormals: [torch.ones(2, 3), torch.ones(2, 3), torch.ones(2)],
+    dist.MixtureOfDiagNormalsSharedCovariance: [torch.ones(2, 3), torch.ones(3), torch.ones(2)],
+    dist.Multinomial: [2, torch.ones(2)],
+    dist.MultivariateNormal: [torch.ones(2), torch.eye(2)],
+    dist.OneHotCategorical: [torch.ones(2)],
+    dist.RelaxedBernoulli: [1.0, 0.5],
+    dist.RelaxedBernoulliStraightThrough: [1.0, 0.5],
+    dist.RelaxedOneHotCategorical: [1., torch.ones(2)],
+    dist.RelaxedOneHotCategoricalStraightThrough: [1., torch.ones(2)],
+    dist.TransformedDistribution: [dist.Normal(0, 1), torch.distributions.ExpTransform()],
+    dist.Uniform: [0, 1],
+    dist.VonMises3D: [torch.tensor([1., 0., 0.])],
+}
+
+
+@pytest.mark.parametrize('Dist', DISTRIBUTIONS)
+def test_pickle(Dist):
+    if Dist in ARGS:
+        args = ARGS[Dist]
+    else:
+        # Optimistically try to initialize with Dist(1, 1, ..., 1).
+        try:
+            # Python 3.6+
+            spec = list(inspect.signature(Dist.__init__).parameters.values())
+            nargs = sum(1 for p in spec if p.default is p.empty) - 1
+        except AttributeError:
+            # Python 2.6-3.5
+            spec = inspect.getargspec(Dist.__init__)
+            nargs = len(spec.args) - 1 - (len(spec.defaults) if spec.defaults else 0)
+        args = (1,) * nargs
+    try:
+        dist = Dist(*args)
+    except Exception:
+        pytest.skip(reason='cannot construct distribution')
+
+    # Note that pickling torch.Size() requires protocol >= 2
+    serialized = pickle.dumps(dist, pickle.HIGHEST_PROTOCOL)
+    deserialized = pickle.loads(serialized)
+    assert isinstance(deserialized, Dist)