Avoiding an infinite loop in a divide-and-conquer algorithm with lax.cond

[s-a-barnett]

I have implemented the divide-and-conquer algorithm for matrix inversion in JAX as follows:

import numpy as np
import jax.numpy as jnp
from jax import jit
from jax import lax
import numpy.random as npr

@jit
def jax_inverse(B):
    '''
    Return the inverse of B using the divide-and-conquer algorithm.
    len(B) is a power of 2.
    '''
    def true_fun(A):
        Ainv = jnp.linalg.inv(A)
        return Ainv
    def false_fun(A):
        # subdivide A into equal quarters
        m = len(A) // 2
        A11 = A[:m, :m]
        A12 = A[:m, m:]
        A21 = A[m:, :m]
        A22 = A[m:, m:]
        # recursive call 1
        A11inv = jax_inverse(A11)
        # recursive call 2
        AS = A22 - (A21 @ A11inv @ A12)
        ASinv = jax_inverse(AS)
        # compute Ainv as a result of two recursive calls
        Ainv = jnp.zeros_like(A)
        Ainv = Ainv.at[:m, :m].set(A11inv + (A11inv @ A12 @ ASinv @ A21 @ A11inv))
        Ainv = Ainv.at[:m, m:].set(-A11inv @ A12 @ ASinv)
        Ainv = Ainv.at[m:, :m].set(-ASinv @ A21 @ A11inv)
        Ainv = Ainv.at[m:, m:].set(ASinv)
        return Ainv
    pred = (len(B) == 2)
    return lax.cond(pred, true_fun, false_fun, B)

if __name__ == '__main__':
    npr.seed(0)
    n = 2 ** 5
    A = npr.randn(n, n)
    Ainv1 = jax_inverse(jnp.array(A))
    Ainv2 = np.linalg.inv(A)
    print(np.linalg.norm(Ainv1 - Ainv2))

Disabling JIT, this works as predicted. However, as written above, we exceed the max recursion depth due as lax.cond traces true_fun even when pred is False. Is there a workaround for this that still uses lax.cond?

[jakevdp]

You cannot recurse on dynamic conditions in JAX due to this issue (both branches being traced). But fortunately in this case the condition is static: i.e. you're branching on the length of the array. So you can replace this:

return lax.cond(pred, true_fun, false_fun, B)

with this:

return true_fun(B) if pred else false_fun(B)

And your code will execute under JIT. Note that with this approach, the entire sequence of operations will be inlined before being sent to the compiler, so it may or may not be a suitable solution to your problem.

[s-a-barnett]

Thanks! I already have a functioning implementation that basically works like this, but I was wondering whether the use of if-then rather than lax.cond means that the overall wall-clock time is slower, possibly slow enough to undermine the use of divide-and-conquer to begin with. Is that what you're elaborating in the last sentence?