Repository: pyro-ppl/numpyro

Issue:
Build wrappers for lax primitives to disable tracing in debug mode
Even with the `disable_jit` flag, it is inconvenient to debug our existing code since we have many utilities that use these lax primitives which need to be manually rewritten each time we want to debug something. Lets build light wrappers around these to make debugging easy. 

We should also move this upstream to JAX if we find it useful.


Discussion:

Resolving patch:
diff --git a/numpyro/hmc_util.py b/numpyro/hmc_util.py
--- a/numpyro/hmc_util.py
+++ b/numpyro/hmc_util.py
@@ -1,11 +1,11 @@
 import jax.numpy as np
-from jax import grad, jit, lax, partial, random, value_and_grad
+from jax import grad, jit, partial, random, value_and_grad
 from jax.flatten_util import ravel_pytree
 from jax.ops import index_update
 from jax.scipy.special import expit
 from jax.tree_util import tree_multimap
 
-from numpyro.util import laxtuple
+from numpyro.util import cond, laxtuple, while_loop
 
 AdaptWindow = laxtuple("AdaptWindow", ["start", "end"])
 AdaptState = laxtuple("AdaptState", ["step_size", "inverse_mass_matrix", "ss_state", "mm_state", "window_idx"])
@@ -174,7 +174,7 @@ def _body_fn(state):
     def _cond_fn(state):
         return (state[1] == 0) | (state[1] == state[2])
 
-    step_size, _, _ = lax.while_loop(_cond_fn, _body_fn, (init_step_size, 0, 0))
+    step_size, _, _ = while_loop(_cond_fn, _body_fn, (init_step_size, 0, 0))
     return step_size
 
 
@@ -272,15 +272,15 @@ def update_fn(t, accept_prob, z_flat, state):
         # update mass matrix state
         is_middle_window = (0 < window_idx) & (window_idx < (num_windows - 1))
         if adapt_mass_matrix:
-            mm_state = lax.cond(is_middle_window,
-                                (z_flat, mm_state), lambda args: mm_update(*args),
-                                mm_state, lambda x: x)
+            mm_state = cond(is_middle_window,
+                            (z_flat, mm_state), lambda args: mm_update(*args),
+                            mm_state, lambda x: x)
 
         t_at_window_end = t == adaptation_schedule[window_idx, 1]
         window_idx = np.where(t_at_window_end, window_idx + 1, window_idx)
         state = AdaptState(step_size, inverse_mass_matrix, ss_state, mm_state, window_idx)
-        state = lax.cond(t_at_window_end & is_middle_window,
-                         state, _update_at_window_end, state, lambda x: x)
+        state = cond(t_at_window_end & is_middle_window,
+                     state, _update_at_window_end, state, lambda x: x)
         return state
 
     return init_fn, update_fn
@@ -326,7 +326,7 @@ def _combine_tree(current_tree, new_tree, inverse_mass_matrix, going_right, rng,
                   biased_transition, iterative_build):
     # Now we combine the current tree and the new tree. Note that outside
     # leaves of the combined tree are determined by the direction.
-    z_left, r_left, z_left_grad, z_right, r_right, r_right_grad = lax.cond(
+    z_left, r_left, z_left_grad, z_right, r_right, r_right_grad = cond(
         going_right,
         (current_tree, new_tree),
         lambda trees: (trees[0].z_left, trees[0].r_left,
@@ -338,14 +338,14 @@ def _combine_tree(current_tree, new_tree, inverse_mass_matrix, going_right, rng,
                        trees[1].r_right, trees[1].z_right_grad)
     )
 
-    transition_prob = lax.cond(biased_transition,
-                               (current_tree, new_tree),
-                               lambda trees: _biased_transition_kernel(trees[0], trees[1]),
-                               (current_tree, new_tree),
-                               lambda trees: _uniform_transition_kernel(trees[0], trees[1]))
+    transition_prob = cond(biased_transition,
+                           (current_tree, new_tree),
+                           lambda trees: _biased_transition_kernel(trees[0], trees[1]),
+                           (current_tree, new_tree),
+                           lambda trees: _uniform_transition_kernel(trees[0], trees[1]))
 
     transition = random.bernoulli(rng, transition_prob)
-    z_proposal, z_proposal_pe, z_proposal_grad = lax.cond(
+    z_proposal, z_proposal_pe, z_proposal_grad = cond(
         transition,
         new_tree, lambda tree: (tree.z_proposal, tree.z_proposal_pe, tree.z_proposal_grad),
         current_tree, lambda tree: (tree.z_proposal, tree.z_proposal_pe, tree.z_proposal_grad)
@@ -424,11 +424,11 @@ def _build_subtree(depth, vv_update, kinetic_fn, z, r, z_grad, inverse_mass_matr
 
 @jit
 def _get_leaf(tree, going_right):
-    return lax.cond(going_right,
-                    tree,
-                    lambda tree: (tree.z_right, tree.r_right, tree.z_right_grad),
-                    tree,
-                    lambda tree: (tree.z_left, tree.r_left, tree.z_left_grad))
+    return cond(going_right,
+                tree,
+                lambda tree: (tree.z_right, tree.r_right, tree.z_right_grad),
+                tree,
+                lambda tree: (tree.z_left, tree.r_left, tree.z_left_grad))
 
 
 def _double_tree(current_tree, vv_update, kinetic_fn, inverse_mass_matrix, step_size,
@@ -455,14 +455,14 @@ def _double_tree(current_tree, vv_update, kinetic_fn, inverse_mass_matrix, step_
 def _leaf_idx_to_ckpt_idxs(n):
     # computes the number of non-zero bits except the last bit
     # e.g. 6 -> 2, 7 -> 2, 13 -> 2
-    _, idx_max = lax.while_loop(lambda nc: nc[0] > 0,
-                                lambda nc: (nc[0] >> 1, nc[1] + (nc[0] & 1)),
-                                (n >> 1, 0))
+    _, idx_max = while_loop(lambda nc: nc[0] > 0,
+                            lambda nc: (nc[0] >> 1, nc[1] + (nc[0] & 1)),
+                            (n >> 1, 0))
     # computes the number of contiguous last non-zero bits
     # e.g. 6 -> 0, 7 -> 3, 13 -> 1
-    _, num_subtrees = lax.while_loop(lambda nc: (nc[0] & 1) != 0,
-                                     lambda nc: (nc[0] >> 1, nc[1] + 1),
-                                     (n, 0))
+    _, num_subtrees = while_loop(lambda nc: (nc[0] & 1) != 0,
+                                 lambda nc: (nc[0] >> 1, nc[1] + 1),
+                                 (n, 0))
     idx_min = idx_max - num_subtrees + 1
     return idx_min, idx_max
 
@@ -474,9 +474,9 @@ def _body_fn(state):
         # XXX we might not need to unravel here
         return i - 1, _is_turning(inverse_mass_matrix, r_ckpts[i], r, subtree_r_sum)
 
-    _, turning = lax.while_loop(lambda it: (it[0] >= idx_min) & ~it[1],
-                                _body_fn,
-                                (idx_max, False))
+    _, turning = while_loop(lambda it: (it[0] >= idx_min) & ~it[1],
+                            _body_fn,
+                            (idx_max, False))
     return turning
 
 
@@ -504,12 +504,12 @@ def _body_fn(state):
         r, _ = ravel_pytree(new_leaf.r_right)
         r_sum, _ = ravel_pytree(new_tree.r_sum)
         # we update checkpoints when leaf_idx is even
-        r_ckpts, r_sum_ckpts = lax.cond(leaf_idx % 2 == 0,
-                                        (r_ckpts, r_sum_ckpts),
-                                        lambda x: (index_update(x[0], ckpt_idx_max, r),
-                                                   index_update(x[1], ckpt_idx_max, r_sum)),
-                                        (r_ckpts, r_sum_ckpts),
-                                        lambda x: x)
+        r_ckpts, r_sum_ckpts = cond(leaf_idx % 2 == 0,
+                                    (r_ckpts, r_sum_ckpts),
+                                    lambda x: (index_update(x[0], ckpt_idx_max, r),
+                                               index_update(x[1], ckpt_idx_max, r_sum)),
+                                    (r_ckpts, r_sum_ckpts),
+                                    lambda x: x)
 
         turning = _is_iterative_turning(inverse_mass_matrix, r, r_sum, r_ckpts, r_sum_ckpts,
                                         ckpt_idx_min, ckpt_idx_max)
@@ -527,7 +527,7 @@ def _body_fn(state):
                                  dtype=inverse_mass_matrix.dtype)
     r_sum_checkpoints = index_update(r_sum_checkpoints, 0, r_init)
 
-    tree, turning, _, _, _ = lax.while_loop(
+    tree, turning, _, _, _ = while_loop(
         _cond_fn,
         _body_fn,
         (basetree, False, r_checkpoints, r_sum_checkpoints, rng)
@@ -579,8 +579,9 @@ def _body_fn(state):
 
     state = (tree, key)
     if iterative_build:
-        tree, _ = lax.while_loop(_cond_fn, _body_fn, state)
+        tree, _ = while_loop(_cond_fn, _body_fn, state)
     else:
         while _cond_fn(state):
-            tree, _ = state = _body_fn(state)
+            state = _body_fn(state)
+        tree, _ = state
     return tree
diff --git a/numpyro/util.py b/numpyro/util.py
--- a/numpyro/util.py
+++ b/numpyro/util.py
@@ -1,12 +1,14 @@
 import random
 from collections import namedtuple
+from contextlib import contextmanager
 
 import numpy as onp
 
+from jax import lax
 from jax.tree_util import register_pytree_node
 
-
 _DATA_TYPES = {}
+_DISABLE_CONTROL_FLOW_PRIM = False
 
 
 def set_rng_seed(rng_seed):
@@ -36,3 +38,56 @@ def laxtuple(name, fields):
     cls.update = cls._replace
     _DATA_TYPES[key] = cls
     return cls
+
+
+@contextmanager
+def optional(condition, context_manager):
+    """
+    Optionally wrap inside `context_manager` if condition is `True`.
+    """
+    if condition:
+        with context_manager:
+            yield
+    else:
+        yield
+
+
+@contextmanager
+def control_flow_prims_disabled():
+    global _DISABLE_CONTROL_FLOW_PRIM
+    stored_flag = _DISABLE_CONTROL_FLOW_PRIM
+    try:
+        _DISABLE_CONTROL_FLOW_PRIM = True
+        yield
+    finally:
+        _DISABLE_CONTROL_FLOW_PRIM = stored_flag
+
+
+def cond(pred, true_operand, true_fun, false_operand, false_fun):
+    if _DISABLE_CONTROL_FLOW_PRIM:
+        if pred:
+            return true_fun(true_operand)
+        else:
+            return false_fun(false_operand)
+    else:
+        return lax.cond(pred, true_operand, true_fun, false_operand, false_fun)
+
+
+def while_loop(cond_fun, body_fun, init_val):
+    if _DISABLE_CONTROL_FLOW_PRIM:
+        val = init_val
+        while cond_fun(val):
+            val = body_fun(val)
+        return val
+    else:
+        return lax.while_loop(cond_fun, body_fun, init_val)
+
+
+def fori_loop(lower, upper, body_fun, init_val):
+    if _DISABLE_CONTROL_FLOW_PRIM:
+        val = init_val
+        for i in range(lower, upper):
+            val = body_fun(i, val)
+        return val
+    else:
+        return lax.fori_loop(lower, upper, body_fun, init_val)


Test patch:
diff --git a/test/test_hmc_util.py b/test/test_hmc_util.py
--- a/test/test_hmc_util.py
+++ b/test/test_hmc_util.py
@@ -6,7 +6,7 @@
 from numpy.testing import assert_allclose
 
 import jax.numpy as np
-from jax import device_put, grad, jit, lax, random, tree_map
+from jax import device_put, disable_jit, grad, jit, lax, random, tree_map
 
 from numpyro.hmc_util import (
     AdaptWindow,
@@ -20,6 +20,7 @@
     warmup_adapter,
     welford_covariance
 )
+from numpyro.util import control_flow_prims_disabled, fori_loop, optional
 
 logger = logging.getLogger(__name__)
 
@@ -47,27 +48,28 @@ def optimize(f):
 @pytest.mark.parametrize('diagonal', [True, False])
 @pytest.mark.parametrize('regularize', [True, False])
 def test_welford_covariance(jitted, diagonal, regularize):
-    onp.random.seed(0)
-    loc = onp.random.randn(3)
-    a = onp.random.randn(3, 3)
-    target_cov = onp.matmul(a, a.T)
-    x = onp.random.multivariate_normal(loc, target_cov, size=(2000,))
-    x = device_put(x)
-
-    def get_cov(x):
-        wc_init, wc_update, wc_final = welford_covariance(diagonal=diagonal)
-        wc_state = wc_init(3)
-        wc_state = lax.fori_loop(0, 2000, lambda i, val: wc_update(x[i], val), wc_state)
-        cov = wc_final(wc_state, regularize=regularize)
-        return cov
-
-    fn = jit(get_cov) if jitted else get_cov
-    cov = fn(x)
-
-    if diagonal:
-        assert_allclose(cov, np.diagonal(target_cov), rtol=0.06)
-    else:
-        assert_allclose(cov, target_cov, rtol=0.06)
+    with optional(jitted, disable_jit()), optional(jitted, control_flow_prims_disabled()):
+        onp.random.seed(0)
+        loc = onp.random.randn(3)
+        a = onp.random.randn(3, 3)
+        target_cov = onp.matmul(a, a.T)
+        x = onp.random.multivariate_normal(loc, target_cov, size=(2000,))
+        x = device_put(x)
+
+        @jit
+        def get_cov(x):
+            wc_init, wc_update, wc_final = welford_covariance(diagonal=diagonal)
+            wc_state = wc_init(3)
+            wc_state = fori_loop(0, 2000, lambda i, val: wc_update(x[i], val), wc_state)
+            cov = wc_final(wc_state, regularize=regularize)
+            return cov
+
+        cov = get_cov(x)
+
+        if diagonal:
+            assert_allclose(cov, np.diagonal(target_cov), rtol=0.06)
+        else:
+            assert_allclose(cov, target_cov, rtol=0.06)
 
 
 ########################################