Repository: pyro-ppl/numpyro

Issue:
Investigate why adaptation does not work well with default HMC parameters
For many models that seem to work well in Pyro with the default values for `trajectory_length` and `step_size` (see `test_mcmc.py`), the behavior in numpyro can be finicky in that either HMC / NUTS is too slow or we get wrong results, despite our tests running for many more steps than in Pyro. 


Discussion:

Resolving patch:
diff --git a/numpyro/mcmc.py b/numpyro/mcmc.py
--- a/numpyro/mcmc.py
+++ b/numpyro/mcmc.py
@@ -41,28 +41,23 @@ def hmc(potential_fn, kinetic_fn=None, algo='NUTS'):
     if kinetic_fn is None:
         kinetic_fn = _euclidean_ke
     vv_init, vv_update = velocity_verlet(potential_fn, kinetic_fn)
-    trajectory_length = None
+    trajectory_len = None
     momentum_generator = None
     wa_update = None
 
     def init_kernel(init_samples,
                     num_warmup_steps,
                     step_size=1.0,
-                    num_steps=None,
                     adapt_step_size=True,
                     adapt_mass_matrix=True,
                     diag_mass=True,
                     target_accept_prob=0.8,
+                    trajectory_length=2*math.pi,
                     run_warmup=True,
                     rng=PRNGKey(0)):
         step_size = float(step_size)
-        nonlocal trajectory_length, momentum_generator, wa_update
-
-        if num_steps is None:
-            trajectory_length = 2 * math.pi
-        else:
-            trajectory_length = num_steps * step_size
-
+        nonlocal momentum_generator, wa_update, trajectory_len
+        trajectory_len = float(trajectory_length)
         z = init_samples
         z_flat, unravel_fn = ravel_pytree(z)
         momentum_generator = partial(_sample_momentum, unravel_fn)
@@ -99,7 +94,7 @@ def warmup_update(t, states):
         return hmc_state, wa_state
 
     def _hmc_next(step_size, inverse_mass_matrix, vv_state, rng):
-        num_steps = _get_num_steps(step_size, trajectory_length)
+        num_steps = _get_num_steps(step_size, trajectory_len)
         vv_state_new = fori_loop(0, num_steps,
                                  lambda i, val: vv_update(step_size, inverse_mass_matrix, val),
                                  vv_state)


Test patch:
diff --git a/test/test_mcmc.py b/test/test_mcmc.py
--- a/test/test_mcmc.py
+++ b/test/test_mcmc.py
@@ -23,7 +23,7 @@ def potential_fn(z):
     init_kernel, sample_kernel = hmc(potential_fn, algo=algo)
     init_samples = np.array(0.)
     hmc_state = init_kernel(init_samples,
-                            num_steps=10,
+                            trajectory_length=10,
                             num_warmup_steps=warmup_steps)
     hmc_states = tscan(lambda state, i: sample_kernel(state), hmc_state, np.arange(num_samples),
                        transform=lambda x: x.z)
@@ -48,8 +48,7 @@ def model(labels):
     init_params, potential_fn, transform_fn = initialize_model(random.PRNGKey(2), model, (labels,), {})
     init_kernel, sample_kernel = hmc(potential_fn, algo=algo)
     hmc_state = init_kernel(init_params,
-                            step_size=0.1,
-                            num_steps=15,
+                            trajectory_length=10,
                             num_warmup_steps=warmup_steps)
     hmc_states = tscan(lambda state, i: sample_kernel(state), hmc_state, np.arange(num_samples),
                        transform=lambda x: transform_fn(x.z))
@@ -73,7 +72,7 @@ def model(data):
     init_kernel, sample_kernel = hmc(potential_fn, algo=algo)
     hmc_state = init_kernel(init_params,
                             step_size=0.1,
-                            num_steps=15,
+                            trajectory_length=1.,
                             num_warmup_steps=warmup_steps)
     hmc_states = tscan(lambda state, i: sample_kernel(state), hmc_state, np.arange(num_samples),
                        transform=lambda x: transform_fn(x.z))
@@ -97,7 +96,7 @@ def model(data):
     init_kernel, sample_kernel = hmc(potential_fn, algo=algo)
     hmc_state = init_kernel(init_params,
                             step_size=0.1,
-                            num_steps=15,
+                            trajectory_length=2.,
                             num_warmup_steps=warmup_steps)
     hmc_states = tscan(lambda state, i: sample_kernel(state), hmc_state, np.arange(num_samples),
                        transform=lambda x: transform_fn(x.z))