Repository: pyro-ppl/pyro

Issue:
Jitting HMC does not work for models with `param` statements
The following repro script shows the issue
```python
import torch
import pyro
import pyro.distributions as dist
from pyro.infer import MCMC, NUTS

def model():
    x = pyro.param("x", torch.tensor(1.))
    pyro.sample("y", dist.Normal(x, 1))

mcmc = MCMC(NUTS(model, jit_compile=True), 10)
mcmc.run()
mcmc.summary()
```
raises the error `RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient`.


Discussion:
An easy workaround is to hide the constant by a `lambda`:
```py
import torch

import pyro
import pyro.distributions as dist
from pyro.infer import MCMC, NUTS

def model():
    x = pyro.param("x", lambda: torch.tensor(1.))
    pyro.sample("y", dist.Normal(x, 1))

model()  # <------ initialize param store

mcmc = MCMC(NUTS(model, jit_compile=True), 10)
mcmc.run()
mcmc.summary()
```
Hmm, I still get the issue. Are you using PyTorch `nightly`?
I'm using PyTorch release. Did you make sure to initialize the param store before inference?
```py
model()  # initialize param store
```
Yes, I used your code and still observed the issue. I think this is the same issue as [this one](https://github.com/pytorch/pytorch/issues/17583) and PyTorch won't support it. IIUC we can resolve the issue by replaying `pyro.param` sites by constant values (using `poutine.substitute` e.g.) in the computation of `potential_fn`.
The JIT tracer expects any constants in the trace to have `requires_grad=False` as expected. Some of our examples and tests for HMC use `model.requires_grad_(False)` if the model is a `PyroModule` or `nn.Module`. 

I think we could provide a `param_store._requires_grad(requires_grad=True)` method which will allow us to use the param store for simple key value lookup that will be useful in such situations (i.e. when model is not a PyroModule). 

```python
def requires_grad_(self, requires_grad=True):
    for v in self._params.values():
        v.requires_grad_(requires_grad)
```



@neerajprad's solution works like a charm!! :D (works with multi-chains too) I think we can call this method automatically at the beginning and the ending of `MCMC.run` method.

Resolving patch:
diff --git a/pyro/infer/mcmc/util.py b/pyro/infer/mcmc/util.py
--- a/pyro/infer/mcmc/util.py
+++ b/pyro/infer/mcmc/util.py
@@ -274,6 +274,12 @@ def _potential_fn_jit(self, skip_jit_warnings, jit_options, params):
             return self._compiled_fn(*vals)
 
         with pyro.validation_enabled(False):
+            tmp = []
+            for _, v in pyro.get_param_store().named_parameters():
+                if v.requires_grad:
+                    v.requires_grad_(False)
+                    tmp.append(v)
+
             def _pe_jit(*zi):
                 params = dict(zip(names, zi))
                 return self._potential_fn(params)
@@ -281,6 +287,9 @@ def _pe_jit(*zi):
             if skip_jit_warnings:
                 _pe_jit = ignore_jit_warnings()(_pe_jit)
             self._compiled_fn = torch.jit.trace(_pe_jit, vals, **jit_options)
+
+            for v in tmp:
+                v.requires_grad_(True)
             return self._compiled_fn(*vals)
 
     def get_potential_fn(self, jit_compile=False, skip_jit_warnings=True, jit_options=None):


Test patch:
diff --git a/tests/infer/mcmc/test_mcmc_util.py b/tests/infer/mcmc/test_mcmc_util.py
--- a/tests/infer/mcmc/test_mcmc_util.py
+++ b/tests/infer/mcmc/test_mcmc_util.py
@@ -54,3 +54,17 @@ def test_predictive(num_samples, parallel):
 
     # check sample mean
     assert_close(predictive_samples["obs"].reshape([-1, 5]).mean(0), true_probs, rtol=0.1)
+
+
+def model_with_param():
+    x = pyro.param("x", torch.tensor(1.))
+    pyro.sample("y", dist.Normal(x, 1))
+
+
+@pytest.mark.parametrize("jit_compile", [False, True])
+@pytest.mark.parametrize("num_chains", [1, 2])
+@pytest.mark.filterwarnings("ignore:num_chains")
+def test_model_with_param(jit_compile, num_chains):
+    kernel = NUTS(model_with_param, jit_compile=jit_compile, ignore_jit_warnings=True)
+    mcmc = MCMC(kernel, 10, num_chains=num_chains, mp_context="spawn")
+    mcmc.run()