Repository: pyro-ppl/numpyro

Issue:
Support forward mode differentiation in HMC
Currently, we use reverse mode differentiation. But forward-mode is also useful in case users' data shapes change and they don't want to recompile the model.


Discussion:

Resolving patch:
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -150,6 +150,7 @@ def hmc(potential_fn=None, potential_fn_gen=None, kinetic_fn=None, algo='NUTS'):
     max_treedepth = None
     wa_update = None
     wa_steps = None
+    forward_mode_ad = False
     max_delta_energy = 1000.
     if algo not in {'HMC', 'NUTS'}:
         raise ValueError('`algo` must be one of `HMC` or `NUTS`.')
@@ -165,6 +166,7 @@ def init_kernel(init_params,
                     trajectory_length=2*math.pi,
                     max_tree_depth=10,
                     find_heuristic_step_size=False,
+                    forward_mode_differentiation=False,
                     model_args=(),
                     model_kwargs=None,
                     rng_key=random.PRNGKey(0)):
@@ -203,7 +205,8 @@ def init_kernel(init_params,
 
         """
         step_size = lax.convert_element_type(step_size, canonicalize_dtype(jnp.float64))
-        nonlocal wa_update, trajectory_len, max_treedepth, vv_update, wa_steps
+        nonlocal wa_update, trajectory_len, max_treedepth, vv_update, wa_steps, forward_mode_ad
+        forward_mode_ad = forward_mode_differentiation
         wa_steps = num_warmup
         trajectory_len = trajectory_length
         max_treedepth = max_tree_depth
@@ -239,7 +242,7 @@ def init_kernel(init_params,
                            inverse_mass_matrix=inverse_mass_matrix,
                            mass_matrix_size=jnp.size(ravel_pytree(z)[0]))
         r = momentum_generator(z, wa_state.mass_matrix_sqrt, rng_key_momentum)
-        vv_init, vv_update = velocity_verlet(pe_fn, kinetic_fn)
+        vv_init, vv_update = velocity_verlet(pe_fn, kinetic_fn, forward_mode_ad)
         vv_state = vv_init(z, r, potential_energy=pe, z_grad=z_grad)
         energy = kinetic_fn(wa_state.inverse_mass_matrix, vv_state.r)
         hmc_state = HMCState(0, vv_state.z, vv_state.z_grad, vv_state.potential_energy, energy,
@@ -249,9 +252,9 @@ def init_kernel(init_params,
     def _hmc_next(step_size, inverse_mass_matrix, vv_state,
                   model_args, model_kwargs, rng_key):
         if potential_fn_gen:
-            nonlocal vv_update
+            nonlocal vv_update, forward_mode_ad
             pe_fn = potential_fn_gen(*model_args, **model_kwargs)
-            _, vv_update = velocity_verlet(pe_fn, kinetic_fn)
+            _, vv_update = velocity_verlet(pe_fn, kinetic_fn, forward_mode_ad)
 
         num_steps = _get_num_steps(step_size, trajectory_len)
         vv_state_new = fori_loop(0, num_steps,
@@ -272,9 +275,9 @@ def _hmc_next(step_size, inverse_mass_matrix, vv_state,
     def _nuts_next(step_size, inverse_mass_matrix, vv_state,
                    model_args, model_kwargs, rng_key):
         if potential_fn_gen:
-            nonlocal vv_update
+            nonlocal vv_update, forward_mode_ad
             pe_fn = potential_fn_gen(*model_args, **model_kwargs)
-            _, vv_update = velocity_verlet(pe_fn, kinetic_fn)
+            _, vv_update = velocity_verlet(pe_fn, kinetic_fn, forward_mode_ad)
 
         binary_tree = build_tree(vv_update, kinetic_fn, vv_state,
                                  inverse_mass_matrix, step_size, rng_key,
@@ -372,6 +375,13 @@ class HMC(MCMCKernel):
         See :ref:`init_strategy` section for available functions.
     :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 bool forward_mode_differentiation: whether to use forward-mode differentiation
+        or reverse-mode differentiation. By default, we use reverse mode but the forward
+        mode can be useful in some cases to improve the performance. In addition, some
+        control flow utility on JAX such as `jax.lax.while_loop` or `jax.lax.fori_loop`
+        only supports forward-mode differentiation. See
+        `JAX's The Autodiff Cookbook <https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html>`_
+        for more information.
     """
     def __init__(self,
                  model=None,
@@ -384,7 +394,8 @@ def __init__(self,
                  target_accept_prob=0.8,
                  trajectory_length=2 * math.pi,
                  init_strategy=init_to_uniform,
-                 find_heuristic_step_size=False):
+                 find_heuristic_step_size=False,
+                 forward_mode_differentiation=False):
         if not (model is None) ^ (potential_fn is None):
             raise ValueError('Only one of `model` or `potential_fn` must be specified.')
         self._model = model
@@ -400,6 +411,7 @@ def __init__(self,
         self._max_tree_depth = 10
         self._init_strategy = init_strategy
         self._find_heuristic_step_size = find_heuristic_step_size
+        self._forward_mode_differentiation = forward_mode_differentiation
         # Set on first call to init
         self._init_fn = None
         self._potential_fn_gen = None
@@ -414,7 +426,8 @@ def _init_state(self, rng_key, model_args, model_kwargs, init_params):
                 dynamic_args=True,
                 init_strategy=self._init_strategy,
                 model_args=model_args,
-                model_kwargs=model_kwargs)
+                model_kwargs=model_kwargs,
+                forward_mode_differentiation=self._forward_mode_differentiation)
             if self._init_fn is None:
                 self._init_fn, self._sample_fn = hmc(potential_fn_gen=potential_fn,
                                                      kinetic_fn=self._kinetic_fn,
@@ -468,6 +481,7 @@ def init(self, rng_key, num_warmup, init_params=None, model_args=(), model_kwarg
             trajectory_length=self._trajectory_length,
             max_tree_depth=self._max_tree_depth,
             find_heuristic_step_size=self._find_heuristic_step_size,
+            forward_mode_differentiation=self._forward_mode_differentiation,
             model_args=model_args,
             model_kwargs=model_kwargs,
             rng_key=rng_key,
@@ -544,6 +558,13 @@ class NUTS(HMC):
         See :ref:`init_strategy` section for available functions.
     :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 bool forward_mode_differentiation: whether to use forward-mode differentiation
+        or reverse-mode differentiation. By default, we use reverse mode but the forward
+        mode can be useful in some cases to improve the performance. In addition, some
+        control flow utility on JAX such as `jax.lax.while_loop` or `jax.lax.fori_loop`
+        only supports forward-mode differentiation. See
+        `JAX's The Autodiff Cookbook <https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html>`_
+        for more information.
     """
     def __init__(self,
                  model=None,
@@ -557,13 +578,15 @@ def __init__(self,
                  trajectory_length=None,
                  max_tree_depth=10,
                  init_strategy=init_to_uniform,
-                 find_heuristic_step_size=False):
+                 find_heuristic_step_size=False,
+                 forward_mode_differentiation=False):
         super(NUTS, self).__init__(potential_fn=potential_fn, model=model, kinetic_fn=kinetic_fn,
                                    step_size=step_size, adapt_step_size=adapt_step_size,
                                    adapt_mass_matrix=adapt_mass_matrix, dense_mass=dense_mass,
                                    target_accept_prob=target_accept_prob,
                                    trajectory_length=trajectory_length,
                                    init_strategy=init_strategy,
-                                   find_heuristic_step_size=find_heuristic_step_size)
+                                   find_heuristic_step_size=find_heuristic_step_size,
+                                   forward_mode_differentiation=forward_mode_differentiation)
         self._max_tree_depth = max_tree_depth
         self._algo = 'NUTS'
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
@@ -3,7 +3,7 @@
 
 from collections import namedtuple
 
-from jax import grad, random, value_and_grad, vmap
+from jax import jacfwd, grad, random, value_and_grad, vmap
 from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 from jax.ops import index_update
@@ -170,6 +170,13 @@ def final_fn(state, regularize=False):
     return init_fn, update_fn, final_fn
 
 
+def _value_and_grad(f, x, forward_mode_differentiation=False):
+    if forward_mode_differentiation:
+        return f(x), jacfwd(f)(x)
+    else:
+        return value_and_grad(f)(x)
+
+
 def _kinetic_grad(kinetic_fn, inverse_mass_matrix, r):
     if hasattr(kinetic_fn, "_kinetic_grad"):
         return kinetic_fn._kinetic_grad(inverse_mass_matrix, r)
@@ -177,7 +184,7 @@ def _kinetic_grad(kinetic_fn, inverse_mass_matrix, r):
         return grad(kinetic_fn, argnums=1)(inverse_mass_matrix, r)
 
 
-def velocity_verlet(potential_fn, kinetic_fn):
+def velocity_verlet(potential_fn, kinetic_fn, forward_mode_differentiation=False):
     r"""
     Second order symplectic integrator that uses the velocity verlet algorithm
     for position `z` and momentum `r`.
@@ -198,7 +205,7 @@ def init_fn(z, r, potential_energy=None, z_grad=None):
         :return: initial state for the integrator.
         """
         if potential_energy is None or z_grad is None:
-            potential_energy, z_grad = value_and_grad(potential_fn)(z)
+            potential_energy, z_grad = _value_and_grad(potential_fn, z, forward_mode_differentiation)
         return IntegratorState(z, r, potential_energy, z_grad)
 
     def update_fn(step_size, inverse_mass_matrix, state):
@@ -213,7 +220,7 @@ def update_fn(step_size, inverse_mass_matrix, state):
         r = tree_multimap(lambda r, z_grad: r - 0.5 * step_size * z_grad, r, z_grad)  # r(n+1/2)
         r_grad = _kinetic_grad(kinetic_fn, inverse_mass_matrix, r)
         z = tree_multimap(lambda z, r_grad: z + step_size * r_grad, z, r_grad)  # z(n+1)
-        potential_energy, z_grad = value_and_grad(potential_fn)(z)
+        potential_energy, z_grad = _value_and_grad(potential_fn, z, forward_mode_differentiation)
         r = tree_multimap(lambda r, z_grad: r - 0.5 * step_size * z_grad, r, z_grad)  # r(n+1)
         return IntegratorState(z, r, potential_energy, z_grad)
 
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -7,7 +7,7 @@
 
 import numpy as np
 
-from jax import device_get, lax, random, value_and_grad
+from jax import device_get, jacfwd, lax, random, value_and_grad
 from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 
@@ -176,7 +176,8 @@ def find_valid_initial_params(rng_key, model,
                               enum=False,
                               model_args=(),
                               model_kwargs=None,
-                              prototype_params=None):
+                              prototype_params=None,
+                              forward_mode_differentiation=False):
     """
     (EXPERIMENTAL INTERFACE) Given a model with Pyro primitives, returns an initial
     valid unconstrained value for all the parameters. This function also returns
@@ -240,7 +241,11 @@ def body_fn(state):
                     key, subkey = random.split(key)
 
         potential_fn = partial(potential_energy, model, model_args, model_kwargs, enum=enum)
-        pe, z_grad = value_and_grad(potential_fn)(params)
+        if forward_mode_differentiation:
+            pe = potential_fn(params)
+            z_grad = jacfwd(potential_fn)(params)
+        else:
+            pe, z_grad = value_and_grad(potential_fn)(params)
         z_grad_flat = ravel_pytree(z_grad)[0]
         is_valid = jnp.isfinite(pe) & jnp.all(jnp.isfinite(z_grad_flat))
         return i + 1, key, (params, pe, z_grad), is_valid
@@ -384,7 +389,8 @@ def initialize_model(rng_key, model,
                      init_strategy=init_to_uniform,
                      dynamic_args=False,
                      model_args=(),
-                     model_kwargs=None):
+                     model_kwargs=None,
+                     forward_mode_differentiation=False):
     """
     (EXPERIMENTAL INTERFACE) Helper function that calls :func:`~numpyro.infer.util.get_potential_fn`
     and :func:`~numpyro.infer.util.find_valid_initial_params` under the hood
@@ -402,6 +408,13 @@ def initialize_model(rng_key, model,
         `potential_fn` and `constraints_fn` callables, respectively.
     :param tuple model_args: args provided to the model.
     :param dict model_kwargs: kwargs provided to the model.
+    :param bool forward_mode_differentiation: whether to use forward-mode differentiation
+        or reverse-mode differentiation. By default, we use reverse mode but the forward
+        mode can be useful in some cases to improve the performance. In addition, some
+        control flow utility on JAX such as `jax.lax.while_loop` or `jax.lax.fori_loop`
+        only supports forward-mode differentiation. See
+        `JAX's The Autodiff Cookbook <https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html>`_
+        for more information.
     :return: a namedtupe `ModelInfo` which contains the fields
         (`param_info`, `potential_fn`, `postprocess_fn`, `model_trace`), where
         `param_info` is a namedtuple `ParamInfo` containing values from the prior
@@ -446,12 +459,14 @@ def initialize_model(rng_key, model,
         unconstrained_values = transform_fn(inv_transforms, init_values, invert=True)
         init_strategy = _init_to_unconstrained_value(values=unconstrained_values)
     prototype_params = transform_fn(inv_transforms, constrained_values, invert=True)
-    (init_params, pe, grad), is_valid = find_valid_initial_params(rng_key, model,
-                                                                  init_strategy=init_strategy,
-                                                                  enum=has_enumerate_support,
-                                                                  model_args=model_args,
-                                                                  model_kwargs=model_kwargs,
-                                                                  prototype_params=prototype_params)
+    (init_params, pe, grad), is_valid = find_valid_initial_params(
+        rng_key, model,
+        init_strategy=init_strategy,
+        enum=has_enumerate_support,
+        model_args=model_args,
+        model_kwargs=model_kwargs,
+        prototype_params=prototype_params,
+        forward_mode_differentiation=forward_mode_differentiation)
 
     if not_jax_tracer(is_valid):
         if device_get(~jnp.all(is_valid)):


Test patch:
diff --git a/test/test_mcmc.py b/test/test_mcmc.py
--- a/test/test_mcmc.py
+++ b/test/test_mcmc.py
@@ -7,7 +7,7 @@
 from numpy.testing import assert_allclose
 import pytest
 
-from jax import device_get, jit, pmap, random, vmap
+from jax import device_get, jit, lax, pmap, random, vmap
 from jax.lib import xla_bridge
 import jax.numpy as jnp
 from jax.scipy.special import logit
@@ -101,7 +101,8 @@ def model(labels):
         assert samples['coefs'].dtype == jnp.float64
 
 
-def test_uniform_normal():
+@pytest.mark.parametrize("forward_mode_differentiation", [True, False])
+def test_uniform_normal(forward_mode_differentiation):
     true_coef = 0.9
     num_warmup, num_samples = 1000, 1000
 
@@ -112,7 +113,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, forward_mode_differentiation=forward_mode_differentiation)
     mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
     mcmc.warmup(random.PRNGKey(2), data, collect_warmup=True)
     assert mcmc.post_warmup_state is not None
@@ -669,3 +670,14 @@ def model():
     mcmc.run(random.PRNGKey(0))
     # because event_shape of x is (1,), x should only take value 1
     assert_allclose(mcmc.get_samples()["x"], jnp.ones((num_samples,) + batch_shape + (1,)))
+
+
+def test_forward_mode_differentiation():
+    def model():
+        x = numpyro.sample("x", dist.Normal(0, 1))
+        y = lax.while_loop(lambda x: x < 10, lambda x: x + 1, x)
+        numpyro.sample("obs", dist.Normal(y, 1), obs=1.)
+
+    # this fails in reverse mode
+    mcmc = MCMC(NUTS(model, forward_mode_differentiation=True), 10, 10)
+    mcmc.run(random.PRNGKey(0))