Repository: pyro-ppl/numpyro

Issue:
Allow to set an additional max tree depth during warmup phase
This could be useful (needs some experiments though) in cases NUTS trajectories are long and it is slow to collect samples to estimate mass matrix during warmup phase. The side effect can be warmup samples are more correlated, so the estimated mass matrix might be not good.


Discussion:
would this be simple to implement? maybe `max_tree_depth` could be an `int` (default behavior) or a `tuple` of `int`s specifying `(max_tree_depth_warmup, max_tree_depth_post_warmup)` 

Resolving patch:
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -247,7 +247,9 @@ def init_kernel(
         :param float trajectory_length: Length of a MCMC trajectory for HMC. Default
             value is :math:`2\\pi`.
         :param int max_tree_depth: Max depth of the binary tree created during the doubling
-            scheme of NUTS sampler. Defaults to 10.
+            scheme of NUTS sampler. Defaults to 10. This argument also accepts a tuple of
+            integers `(d1, d2)`, where `d1` is the max tree depth during warmup phase and
+            `d2` is the max tree depth during post warmup phase.
         :param bool find_heuristic_step_size: whether to a heuristic function to adjust the
             step size at the beginning of each adaptation window. Defaults to False.
         :param tuple model_args: Model arguments if `potential_fn_gen` is specified.
@@ -263,7 +265,11 @@ def init_kernel(
         nonlocal wa_update, max_treedepth, vv_update, wa_steps, forward_mode_ad
         forward_mode_ad = forward_mode_differentiation
         wa_steps = num_warmup
-        max_treedepth = max_tree_depth
+        max_treedepth = (
+            max_tree_depth
+            if isinstance(max_tree_depth, tuple)
+            else (max_tree_depth, max_tree_depth)
+        )
         if isinstance(init_params, ParamInfo):
             z, pe, z_grad = init_params
         else:
@@ -376,7 +382,7 @@ def _nuts_next(
         model_args,
         model_kwargs,
         rng_key,
-        trajectory_length,
+        max_treedepth_current,
     ):
         if potential_fn_gen:
             nonlocal vv_update, forward_mode_ad
@@ -391,7 +397,7 @@ def _nuts_next(
             step_size,
             rng_key,
             max_delta_energy=max_delta_energy,
-            max_tree_depth=max_treedepth,
+            max_tree_depth=(max_treedepth_current, max(max_treedepth)),
         )
         accept_prob = binary_tree.sum_accept_probs / binary_tree.num_proposals
         num_steps = binary_tree.num_proposals
@@ -437,6 +443,12 @@ def sample_kernel(hmc_state, model_args=(), model_kwargs=None):
         vv_state = IntegratorState(
             hmc_state.z, r, hmc_state.potential_energy, hmc_state.z_grad
         )
+        if algo == "HMC":
+            hmc_length_args = (hmc_state.trajectory_length,)
+        else:
+            hmc_length_args = (
+                jnp.where(hmc_state.i < wa_steps, max_treedepth[0], max_treedepth[1]),
+            )
         vv_state, energy, num_steps, accept_prob, diverging = _next(
             hmc_state.adapt_state.step_size,
             hmc_state.adapt_state.inverse_mass_matrix,
@@ -444,7 +456,7 @@ def sample_kernel(hmc_state, model_args=(), model_kwargs=None):
             model_args,
             model_kwargs,
             rng_key_transition,
-            hmc_state.trajectory_length,
+            *hmc_length_args,
         )
         # not update adapt_state after warmup phase
         adapt_state = cond(
@@ -800,7 +812,9 @@ class NUTS(HMC):
     :param float trajectory_length: Length of a MCMC trajectory for HMC. This arg has
         no effect in NUTS sampler.
     :param int max_tree_depth: Max depth of the binary tree created during the doubling
-        scheme of NUTS sampler. Defaults to 10.
+        scheme of NUTS sampler. Defaults to 10. This argument also accepts a tuple of
+        integers `(d1, d2)`, where `d1` is the max tree depth during warmup phase and
+        `d2` is the max tree depth during post warmup phase.
     :param callable init_strategy: a per-site initialization function.
         See :ref:`init_strategy` section for available functions.
     :param bool find_heuristic_step_size: whether to a heuristic function to adjust the
diff --git a/numpyro/infer/hmc_util.py b/numpyro/infer/hmc_util.py
--- a/numpyro/infer/hmc_util.py
+++ b/numpyro/infer/hmc_util.py
@@ -1116,9 +1116,17 @@ def build_tree(
         randomness.
     :param float max_delta_energy: A threshold to decide if the new state diverges
         (based on the energy difference) too much from the initial integrator state.
+    :param int max_tree_depth: Max depth of the binary tree created during the doubling
+        scheme of NUTS sampler. Defaults to 10. This argument also accepts a tuple of
+        integers `(d1, d2)`, where `d1` is the max tree depth at the current MCMC
+        step and `d2` is the global max tree depth for all MCMC steps.
     :return: information of the tree.
     :rtype: :data:`TreeInfo`
     """
+    if isinstance(max_tree_depth, tuple):
+        max_tree_depth_current, max_tree_depth = max_tree_depth
+    else:
+        max_tree_depth_current = max_tree_depth
     z, r, potential_energy, z_grad = verlet_state
     energy_current = potential_energy + kinetic_fn(inverse_mass_matrix, r)
     latent_size = jnp.size(ravel_pytree(r)[0])
@@ -1147,7 +1155,7 @@ def build_tree(
 
     def _cond_fn(state):
         tree, _ = state
-        return (tree.depth < max_tree_depth) & ~tree.turning & ~tree.diverging
+        return (tree.depth < max_tree_depth_current) & ~tree.turning & ~tree.diverging
 
     def _body_fn(state):
         tree, key = state


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
@@ -156,7 +156,8 @@ def model(data):
     assert_allclose(jnp.mean(samples["loc"], 0), true_coef, atol=0.05)
 
 
-def test_improper_normal():
+@pytest.mark.parametrize("max_tree_depth", [10, (5, 10)])
+def test_improper_normal(max_tree_depth):
     true_coef = 0.9
 
     def model(data):
@@ -171,7 +172,7 @@ def model(data):
         numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
 
     data = true_coef + random.normal(random.PRNGKey(0), (1000,))
-    kernel = NUTS(model=model)
+    kernel = NUTS(model=model, max_tree_depth=max_tree_depth)
     mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000)
     mcmc.run(random.PRNGKey(0), data)
     samples = mcmc.get_samples()