Repository: pyro-ppl/pyro

Issue:
[bug] MCMC with potential_fn gives error
### Issue Description

Since pyro v1.4.0, MCMC with a `potential_fn` fails if `transforms=None` (default value) is provided, see code snippet below. It worked fine for v1.3.1.

A work-around is explicitly passing `transforms={}`. However, I could not find a documentation of this API change. Is this intentional or a bug?

### Environment

macOS 10.13.6, python 3.8.2.
torch: 1.6.0
pyro: 1.4.0

### Code Snippet

```python
import torch
import pyro
from pyro.infer import HMC, MCMC

init_params = {"z": torch.tensor(0.)}

def potential_fn(params):
    return params["z"]

# Fails in v1.4.0, works in v.1.3.1
mcmc = MCMC(
    kernel=HMC(potential_fn=potential_fn),
    num_samples=10,
    warmup_steps=10,
    transforms=None,  # None is the default value
    initial_params=init_params)
mcmc.run()  # On v1.4.0 yields TypeError: 'NoneType' object is not subscriptable

# Work-around for v1.4.0
mcmc = MCMC(
    kernel=HMC(potential_fn=potential_fn),
    num_samples=10,
    warmup_steps=10,
    transforms={},
    initial_params=init_params)
mcmc.run()
```

### Analysis

From what I understand, after the chain was run, but **before** `transforms` are applied, `.cleanup` is called [here](https://github.com/pyro-ppl/pyro/blob/2dfa8da0dd400c3712768385d8306848e93dab9a/pyro/infer/mcmc/api.py#L171) which sets the `initial_params=None`.

Since v1.4.0, `self.kernel.setup()` will be called [here](https://github.com/pyro-ppl/pyro/blob/2dfa8da0dd400c3712768385d8306848e93dab9a/pyro/infer/mcmc/api.py#L414) if `self.kernel.transforms is None`. This will then throw an error due to `initial_params=None`. Until pyro v1.3.1, `self.transforms` had simply been set to `{}` if `self.kernels.transforms is None`.


Discussion:

Patch:
diff --git a/pyro/distributions/mixture.py b/pyro/distributions/mixture.py
index 260c3b12..6eb8378c 100644
--- a/pyro/distributions/mixture.py
+++ b/pyro/distributions/mixture.py
@@ -48,12 +48,12 @@ class MaskedMixture(TorchDistribution):
         with pyro.plate("data", len(data)):
             pyro.sample("obs", MaskedMixture(mask, dist1, dist2), obs=data)
 
-    :param torch.Tensor mask: A byte tensor toggling between ``component0``
+    :param torch.Tensor mask: A boolean tensor toggling between ``component0``
         and ``component1``.
     :param pyro.distributions.TorchDistribution component0: a distribution
-        for batch elements ``mask == 0``.
+        for batch elements ``mask == False``.
     :param pyro.distributions.TorchDistribution component1: a distribution
-        for batch elements ``mask == 1``.
+        for batch elements ``mask == True``.
     """
     arg_constraints = {}  # nothing can be constrained
 
diff --git a/pyro/infer/mcmc/api.py b/pyro/infer/mcmc/api.py
index 6169073d..dfa55cf1 100644
--- a/pyro/infer/mcmc/api.py
+++ b/pyro/infer/mcmc/api.py
@@ -26,7 +26,7 @@ import pyro.poutine as poutine
 from pyro.infer.mcmc.hmc import HMC
 from pyro.infer.mcmc.logger import DIAGNOSTIC_MSG, ProgressBar, TqdmHandler, initialize_logger
 from pyro.infer.mcmc.nuts import NUTS
-from pyro.infer.mcmc.util import diagnostics, initialize_model, print_summary
+from pyro.infer.mcmc.util import diagnostics, print_summary
 from pyro.util import optional
 
 MAX_SEED = 2**32 - 1
@@ -408,19 +408,14 @@ class MCMC:
         # If transforms is not explicitly provided, infer automatically using
         # model args, kwargs.
         if self.transforms is None:
-            # Try to initialize kernel.transforms using kernel.setup().
-            if getattr(self.kernel, "transforms", None) is None:
-                warmup_steps = 0
-                self.kernel.setup(warmup_steps, *args, **kwargs)
             # Use `kernel.transforms` when available
             if getattr(self.kernel, "transforms", None) is not None:
                 self.transforms = self.kernel.transforms
             # Else, get transforms from model (e.g. in multiprocessing).
             elif self.kernel.model:
-                _, _, self.transforms, _ = initialize_model(self.kernel.model,
-                                                            model_args=args,
-                                                            model_kwargs=kwargs,
-                                                            initial_params={})
+                warmup_steps = 0
+                self.kernel.setup(warmup_steps, *args, **kwargs)
+                self.transforms = self.kernel.transforms
             # Assign default value
             else:
                 self.transforms = {}


Test patch:
diff --git a/tests/infer/mcmc/test_mcmc_api.py b/tests/infer/mcmc/test_mcmc_api.py
index 30ea24e9..a577da9d 100644
--- a/tests/infer/mcmc/test_mcmc_api.py
+++ b/tests/infer/mcmc/test_mcmc_api.py
@@ -230,3 +230,17 @@ def test_sequential_consistent(monkeypatch):
 
     assert_close(samples1["y"][0], samples2["y"][1])
     assert_close(samples1["y"][1], samples2["y"][0])
+
+
+def test_model_with_potential_fn():
+    init_params = {"z": torch.tensor(0.)}
+
+    def potential_fn(params):
+        return params["z"]
+
+    mcmc = MCMC(
+        kernel=HMC(potential_fn=potential_fn),
+        num_samples=10,
+        warmup_steps=10,
+        initial_params=init_params)
+    mcmc.run()