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:
I think the confusion here is that HMC respects `trajectory_length`. So `step_size=trajectory_length / num_steps`. In your case, trajectory_length is large, so step_size will be large. This is confusing. We can't adapt step size if we have a fixed num_steps. But I guess we need to improve the tqdm bar to report the actual step_size - currently, it reports the input value (which will be changed if num_steps is specified).
We can also raise a warning if `num_steps` is specified and improve the docstrings correspondingly.
Thank you for the quick response.
Is it possible to fix the number of LF steps per sample and tune the
stepsize to the desired acceptance rate?

On Mon, Nov 13, 2023 at 2:09 PM Du Phan ***@***.***> wrote:

> We can also raise a warning if num_steps is specified.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/pyro-ppl/numpyro/issues/1678#issuecomment-1808136015>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AKIPD4HD4BYVPQ5XFJXI3ELYEILYTAVCNFSM6AAAAAA7I4ULF6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQMBYGEZTMMBRGU>
> .
> You are receiving this because you authored the thread.Message ID:
> ***@***.***>
>

It can be supported by allowing setting `trajectory_length=None` or something. Currently, with a fixed trajectory length and fixed number of steps, the step size will also be fixed.

Resolving patch:
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -5,6 +5,7 @@
 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 @@
     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 @@
  - **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_next(
             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 @@ def __init__(
     ):
         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
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -537,7 +537,7 @@ def warmup(
         **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
--- a/setup.py
+++ b/setup.py
@@ -63,7 +63,7 @@
             "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
--- a/test/infer/test_mcmc.py
+++ b/test/infer/test_mcmc.py
@@ -1079,7 +1079,8 @@ def model(data):
         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 model(data):
     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)