O(n)-sized jaxpr takes O(exp(n)) work to obtain

The following code has exponential compilation time in the depth (here passed as 10):

import jax
import jax.lax as lax
import jax.numpy as jnp

def fn(x, depth):
    if depth == 1:
        return x
    else:
        _fn = lambda x: fn(x, depth - 1)
        return lax.cond(jnp.array(True), _fn, lambda x: x, x)

jax.make_jaxpr(jax.vmap(lambda x: fn(x, 10)))(jnp.array([1.]))

The reason is that the condition jnp.array(True) is not vmap’d, which means that two calls to batch_jaxpr are made:

Here: https://github.com/google/jax/blob/d8b7bd54be7252673a211b8f4d06a0b4d72256c1/jax/_src/lax/control_flow.py#L857 and here: https://github.com/google/jax/blob/d8b7bd54be7252673a211b8f4d06a0b4d72256c1/jax/_src/lax/control_flow.py#L861 So that overall 2^depth calls are made to batch_jaxpr.

I’m currently writing something involving a binary tree, but unfortunately this is pretty much a blocker.

[apaszke]

Yeah, adding jax._src.util.cache around batching.batch_jaxpr and adding some tuple casts for its arguments drops the execution time from 80s to milliseconds.

[mattjj]

Hrm really doesn’t seem exponential anymore: