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:
@fehiepsi I suspect this is related to our initialization refactoring, e.g. #2417. I'll take a closer look later, but do you have any immediate ideas?
I think this happens after we put some logics for [epidemiology](https://github.com/pyro-ppl/pyro/pull/2527/files#diff-6075bc4ce62e591e03146eae732344e7R411). I guess we can move that logic to the branch `elif self.kernel.model:`?

Resolving patch:
diff --git a/pyro/infer/mcmc/api.py b/pyro/infer/mcmc/api.py
--- a/pyro/infer/mcmc/api.py
+++ b/pyro/infer/mcmc/api.py
@@ -26,7 +26,7 @@
 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 @@ def model(data):
         # 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
--- a/tests/infer/mcmc/test_mcmc_api.py
+++ b/tests/infer/mcmc/test_mcmc_api.py
@@ -230,3 +230,17 @@ def setup(self, warmup_steps, *args, **kwargs):
 
     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()