Repository: pyro-ppl/pyro
Pull request: Fix pickling of dynamically imported torch distributions
Fixes #1456 

Note that pickling`torch.Size()` requires pickle protocol >= 2.

`Gumbel` currently fails due to its use of an inverse transform. I believe this can be fixed upstream by adding `__getstate__` and `__setstate__` methods to `InverseTransform`.

## Tested

- Added a `pickle.dumps()` test for all Pyro distributions.

Problem statement:
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.

Hints from issue 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?

Commit message: Fix pickling of dynamically imported torch distributions

Patch to review:
diff --git a/pyro/distributions/torch.py b/pyro/distributions/torch.py
index ea82e5504b..ccd575b1ae 100644
--- 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__ = '''
diff --git a/tests/distributions/test_pickle.py b/tests/distributions/test_pickle.py
new file mode 100644
index 0000000000..fe6b63c790
--- /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:
+        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)


Code review comments:
diff_hunk: @@ -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:
path: tests/distributions/test_pickle.py
text: @user1:
You'll need to add `Exception` to make the linter happy.