Compute logprob & grad at the same time

Author: rlouf

blackjax/mcmc/diffusion.py

@@ -11,6 +11,7 @@
 
 class DiffusionState(NamedTuple):
     position: PyTree
+    logprob: float
     logprob_grad: PyTree
 
 
@@ -24,7 +25,7 @@ def overdamped_langevin(logprob_grad_fn):
     """Euler solver for overdamped Langevin diffusion."""
 
     def one_step(rng_key, state: DiffusionState, step_size: float, batch: tuple = ()):
-        position, logprob_grad = state
+        position, _, logprob_grad = state
         noise = generate_gaussian_noise(rng_key, position)
         position = jax.tree_util.tree_multimap(
             lambda p, g, n: p + step_size * g + jnp.sqrt(2 * step_size) * n,
@@ -33,7 +34,7 @@ def one_step(rng_key, state: DiffusionState, step_size: float, batch: tuple = ()
             noise,
         )
 
-        logprob_grad = logprob_grad_fn(position, *batch)
-        return DiffusionState(position, logprob_grad)
+        logprob, logprob_grad = logprob_grad_fn(position, *batch)
+        return DiffusionState(position, logprob, logprob_grad)
 
     return one_step

blackjax/mcmc/mala.py

@@ -21,6 +21,7 @@ class MALAState(NamedTuple):
     """
 
     position: PyTree
+    logprob: float
     logprob_grad: PyTree
 
 
@@ -43,9 +44,9 @@ class MALAInfo(NamedTuple):
 
 
 def init(position: PyTree, logprob_fn: Callable) -> MALAState:
-    grad_fn = jax.grad(logprob_fn)
-    logprob_grad = grad_fn(position)
-    return MALAState(position, logprob_grad)
+    grad_fn = jax.value_and_grad(logprob_fn)
+    logprob, logprob_grad = grad_fn(position)
+    return MALAState(position, logprob, logprob_grad)
 
 
 def kernel():
@@ -72,16 +73,16 @@ def one_step(
         TODO expand the docstring.
 
         """
-        grad_fn = jax.grad(logprob_fn)
+        grad_fn = jax.value_and_grad(logprob_fn)
         integrator = overdamped_langevin(grad_fn)
 
         key_integrator, key_rmh = jax.random.split(rng_key)
 
         new_state = integrator(key_integrator, state, step_size)
 
         delta = (
-            logprob_fn(new_state.position)
-            - logprob_fn(state.position)
+            new_state.logprob
+            - state.logprob
             + transition_probability(new_state, state, step_size)
             - transition_probability(state, new_state, step_size)
         )