Repository: pyro-ppl/numpyro

Issue:
Zero HMC acceptance rate for standard normal
Step size tuning in HMC fails on a standard Gaussian. Acceptance rate is 0, regardless of the step size. The problem goes away with NUTS. 
The problem is shown in this notebook: 
https://colab.research.google.com/drive/1QBrROnhUEAOcP3Q9-pT8AfwFjZdu9y1y?usp=sharing

Discussion:
fehiepsi: I guess the latest working version is `jaxns=2.2.6`. Could you limit the version [here](https://github.com/pyro-ppl/numpyro/blob/master/setup.py#L66) and [here](https://github.com/pyro-ppl/numpyro/blob/master/docs/requirements.txt#L7) 
fehiepsi: Thanks, @yayami3! Could you add a simple test for this?
yayami3: test added. 

@fehiepsi 
pls review!

Patch:
diff --git a/docs/requirements.txt b/docs/requirements.txt
index 36b2fb48..b42df688 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -4,7 +4,7 @@ funsor
 ipython<=8.6.0  # strict the version for https://github.com/ipython/ipython/issues/13845
 jax
 jaxlib
-jaxns>=2.0.1
+jaxns==2.2.6
 Jinja2<3.1
 matplotlib
 multipledispatch
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
index 8aac9e4d..f3b31adf 100644
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -5,6 +5,7 @@ from collections import OrderedDict, namedtuple
 from functools import partial
 import math
 import os
+import warnings
 
 from jax import device_put, lax, random, vmap
 from jax.flatten_util import ravel_pytree
@@ -19,7 +20,12 @@ from numpyro.infer.hmc_util import (
     warmup_adapter,
 )
 from numpyro.infer.mcmc import MCMCKernel
-from numpyro.infer.util import ParamInfo, init_to_uniform, initialize_model
+from numpyro.infer.util import (
+    ParamInfo,
+    find_stack_level,
+    init_to_uniform,
+    initialize_model,
+)
 from numpyro.util import cond, fori_loop, identity, is_prng_key
 
 HMCState = namedtuple(
@@ -54,6 +60,7 @@ A :func:`~collections.namedtuple` consisting of the following fields:
  - **trajectory_length** - The amount of time to run HMC dynamics in each sampling step.
    This field is not used in NUTS.
  - **num_steps** - Number of steps in the Hamiltonian trajectory (for diagnostics).
+   In HMC sampler, `trajectory_length` should be None for step_size to be adapted.
    In NUTS sampler, the tree depth of a trajectory can be computed from this field
    with `tree_depth = np.log2(num_steps).astype(int) + 1`.
  - **accept_prob** - Acceptance probability of the proposal. Note that ``z``
@@ -362,8 +369,10 @@ def hmc(potential_fn=None, potential_fn_gen=None, kinetic_fn=None, algo="NUTS"):
             num_steps = 1
         else:
             num_steps = _get_num_steps(step_size, trajectory_length)
-        # makes sure trajectory length is constant, rather than step_size * num_steps
-        step_size = trajectory_length / num_steps
+
+        if trajectory_length is not None:
+            # makes sure trajectory length is constant, rather than step_size * num_steps
+            step_size = trajectory_length / num_steps
         vv_state_new = fori_loop(
             0,
             num_steps,
@@ -618,6 +627,20 @@ class HMC(MCMCKernel):
     ):
         if not (model is None) ^ (potential_fn is None):
             raise ValueError("Only one of `model` or `potential_fn` must be specified.")
+        if type(self) is HMC:
+            if (num_steps is None) & (trajectory_length is None):
+                raise ValueError(
+                    "At least one of `num_steps` or `trajectory_length` must be specified."
+                )
+            if (
+                adapt_step_size
+                & (num_steps is not None)
+                & (trajectory_length is not None)
+            ):
+                warnings.warn(
+                    "If both `num_steps` and `trajectory_length` are specified step size can't be adapted",
+                    stacklevel=find_stack_level(),
+                )
         self._model = model
         self._potential_fn = potential_fn
         self._kinetic_fn = (
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
index 92f461a7..0b1fb085 100644
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -537,7 +537,7 @@ class MCMC(object):
         **kwargs,
     ):
         """
-        Run the MCMC warmup adaptation phase. After this call, `self.warmup_state` will be set
+        Run the MCMC warmup adaptation phase. After this call, `self.post_warmup_state` will be set
         and the :meth:`run` method will skip the warmup adaptation phase. To run `warmup` again
         for the new data, it is required to run :meth:`warmup` again.
 
diff --git a/setup.py b/setup.py
index a14e56a9..fc346e03 100644
--- a/setup.py
+++ b/setup.py
@@ -63,7 +63,7 @@ setup(
             "flax",
             "funsor>=0.4.1",
             "graphviz",
-            "jaxns>=2.0.1",
+            "jaxns==2.2.6",
             "matplotlib",
             "optax>=0.0.6",
             "pylab-sdk",  # jaxns dependency


Test patch:
diff --git a/test/infer/test_mcmc.py b/test/infer/test_mcmc.py
index 368fde59..04a5bc64 100644
--- a/test/infer/test_mcmc.py
+++ b/test/infer/test_mcmc.py
@@ -1079,7 +1079,8 @@ def test_fixed_num_steps():
         with numpyro.plate("size", np.size(data["y"])):
             numpyro.sample("obs", dist.Normal(w * data["x"] + b, sigma), obs=data["y"])
 
-    hmc_kernel = HMC(model, num_steps=5)
+    with pytest.warns(UserWarning, match="If both"):
+        hmc_kernel = HMC(model, num_steps=5)
     mcmc = MCMC(
         hmc_kernel,
         num_samples=1000,
@@ -1090,3 +1091,35 @@ def test_fixed_num_steps():
     mcmc.run(rng_key, data, extra_fields=("num_steps",))
     num_steps_list = np.array(mcmc.get_extra_fields()["num_steps"])
     assert all(step == 5 for step in num_steps_list)
+
+
+@pytest.mark.parametrize("num_steps", [None, 10])
+def test_none_trajectory_length(num_steps):
+    data = dict()
+    data["x"] = np.random.rand(10)
+    data["y"] = data["x"] + np.random.rand(10) * 0.1
+
+    def model(data):
+        w = numpyro.sample("w", dist.Normal(10, 1))
+        b = numpyro.sample("b", dist.Normal(1, 1))
+        sigma = numpyro.sample("sigma", dist.Gamma(1, 2))
+        with numpyro.plate("size", np.size(data["y"])):
+            numpyro.sample("obs", dist.Normal(w * data["x"] + b, sigma), obs=data["y"])
+
+    if num_steps is None:
+        with pytest.raises(ValueError, match="At least one of"):
+            hmc_kernel = HMC(model, num_steps=num_steps, trajectory_length=None)
+    else:
+        hmc_kernel = HMC(model, num_steps=num_steps, trajectory_length=None)
+
+        mcmc = MCMC(
+            hmc_kernel,
+            num_samples=1000,
+            num_warmup=1000,
+            num_chains=1,
+        )
+        rng_key = random.PRNGKey(0)
+        mcmc.run(rng_key, data)
+        mcmc.run(rng_key, data, extra_fields=("num_steps",))
+        num_steps_list = np.array(mcmc.get_extra_fields()["num_steps"])
+        assert all(step == num_steps for step in num_steps_list)