Execution count for loop + vmap

[f0uriest]

Suppose I have some code like this:

def bar(n):
    x = 0
    def body(i, x):
        x += fun(x)
        return x

    return jax.lax.fori_loop(0, n, body, x)

out = jax.vmap(bar)(np.arange(m))

My question is, how many times will fun be called in total?
If I replaced vmap with an explicit loop, it would be called 0 + 1 + 2 + ... + m-1 = m*(m-1)/2 times, but with the vmap I'm not sure how it gets plumbed down. Because the loop bounds aren't static I think it becomes a while loop, which seems like they get vectorized in such a way that all the loops run for the full m steps (#2947), so the total would be like m**2 calls to fun

[jakevdp]

The Python function fun will be called only once with a tracer in place of x. This will evaluate the effect of fun on any input of the given shape and dtype, and then this lowered code will be executed m times on the device without calling fun again.

For a vmapped fori_loop with static bounds, it will be changed to a while_loop that will perform m iterations on the batched version of fun, again without calling the python function fun again.

For more background on JAX's trace-based execution model, see How To Think In JAX.

[f0uriest]

I guess I meant more along the lines of "how many times will the scalar ops in fun get executed". Or put another way, how does vectorizing over a loop termination condition work. It sounds like the total count would be m**2 (m steps in the loop, and each call to vectorized fun is ~ m ops, though parallelized)

[jakevdp]

fori_loop with a static bound lowers to scan, and fori_loop with a dynamic bound lowers to while_loop. Since you are mapping over n here, it lowers to a while_loop.

When you vmap a while_loop over the bound, it results in a single while_loop over the batched body function. So in the end there will be m iterations over a vmapped call to f with a batch size of m.

Does that answer your question?

[f0uriest]

Yes that's perfect, thanks!