lax.scan gives the same performance as normal for loop?

[dui1234]

Hi,

Based on my previous question #15134 (comment), I added a question related to it but there is no answer so far. I, thus, create the new question specifically for the additional question.

I have the follwing code that do some operations on a tree-like structure, a binary tree:

import jax.numpy as jnp
import jax
from jax import jit, vmap, lax, random
from jax.experimental import checkify
from functools import partial
import time

@jit
def stack_p(params):
    return jax.tree_map(lambda *x: jnp.stack(x), *params)

def construct_tree(size, init_value, alpha: float = 0.6):
    tree_capacity = 1
    while tree_capacity < size: tree_capacity *= 2
    return [jnp.full(shape = (2 * tree_capacity), fill_value=init_value),
            tree_capacity, 
            size, 
            alpha, 
            1.0, 
            0] 

def construct_trees(size, alpha: float = 0.6):
    return [construct_tree(size, 0.0, alpha), construct_tree(size, float("inf"), alpha)]

@partial(jit,static_argnums = (0))
def set_tree_value(function,tree,idx,value):
    idx += tree[1]
    tree[0] = tree[0].at[idx].set(value)
    idx //= 2
    cond_fun = lambda x: x[0]>= 1
    def body_fun(idx_t):
        tree = idx_t[1].at[idx_t[0]].set(function(idx_t[1][2 * idx_t[0]], idx_t[1][2 * idx_t[0] + 1]))
        idx = idx_t[0]//2
        return [idx,tree]
    dmy = lax.while_loop(cond_fun, body_fun, [idx,tree[0]])
    tree[0] = dmy[1]
    return tree

@jit
def set_trees_value(trees):
    trees[0] = set_tree_value(lax.add,trees[0],trees[0][-1],trees[0][-2]**trees[0][3])
    trees[0][-1] = (trees[0][-1] + 1) % trees[0][2]
    trees[1] = set_tree_value(lax.min,trees[1],trees[1][-1],trees[1][-2]**trees[1][3])
    trees[1][-1] = (trees[1][-1] + 1) % trees[1][2]
    return trees

jv_set_trees_value = jit(vmap(set_trees_value))

trees = [construct_trees(1000, 0.2)]
trees_Ns = stack_p(trees)
for i in range(2000): trees_Ns = vmap(set_trees_value)(trees_Ns)

seed = 999
key = random.PRNGKey(seed)
priority_dmy = random.normal(key,(1, 100, 1))
indx = [jnp.sort(random.choice(key,jnp.arange(1000),(100,)))]
tree_dmy = [jax.tree_map(lambda x:x[i], trees_Ns) for i in range(1)]

@jit
def update_priorities2(trees, indices, priorities):
    f2 = lambda x:x>0
    checked_f2 = checkify.checkify(f2, errors=checkify.all_checks)
    err = vmap(checked_f2)(priorities)
    max_prior = jnp.max(lax.max(priorities,trees[0][4]))
    trees[0][-2] = max_prior
    trees[1][-2] = max_prior
    
    for idx, priority in zip(indices, priorities.reshape(-1)):
        trees[0] = set_tree_value(lax.add,trees[0],idx,priority**trees[0][3])
        trees[1] = set_tree_value(lax.min,trees[1],idx,priority**trees[1][3])
    
    return trees 

v_update_tree = vmap(update_priorities2)
j_update_tree = jit(vmap(update_priorities2))

This time used for performing the functuion update_priorities2 along with its vmapped transformed versions are following:

_ = update_priorities2(tree_dmy[0],stack_p(indx)[0],priority_dmy[0])  # compile
%timeit jax.block_until_ready(update_priorities2(tree_dmy[0],stack_p(indx)[0],priority_dmy[0]))
# 1.49 ms ± 244 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

_ = j_update_tree(trees_Ns,stack_p(indx),priority_dmy)  # compile
%timeit jax.block_until_ready(j_update_tree(trees_Ns,stack_p(indx),priority_dmy))
# 312 µs ± 12.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

_ = v_update_tree(trees_Ns,stack_p(indx),priority_dmy)  # compile
%timeit jax.block_until_ready(v_update_tree(trees_Ns,stack_p(indx),priority_dmy))
# 2.96 ms ± 119 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

I, thus, replaced the for loop inside the function with the lax.scan and get the operation times as:

@jit
def update_priorities2_new(trees, indices, priorities):
    f2 = lambda x:x>0
    checked_f2 = checkify.checkify(f2, errors=checkify.all_checks)
    err = vmap(checked_f2)(priorities)
    max_prior = jnp.max(lax.max(priorities,trees[0][4]))
    trees[0][-2] = max_prior
    trees[1][-2] = max_prior

    def dmy_update(trees, idx_pri):
        idx, priority = idx_pri
        trees[0] = set_tree_value(lax.add,trees[0],idx,priority**trees[0][3])
        trees[1] = set_tree_value(lax.min,trees[1],idx,priority**trees[1][3])
        return trees, trees
    
    r,a = lax.scan(dmy_update,init = trees, xs = (indices,priorities.reshape(-1)))
    return r

v_update_tree_new = vmap(update_priorities2_new)
j_update_tree_new = jit(vmap(update_priorities2_new))

_ = update_priorities2_new(tree_dmy[0],stack_p(indx)[0],priority_dmy[0])  # compile
%timeit jax.block_until_ready(update_priorities2(tree_dmy[0],stack_p(indx)[0],priority_dmy[0]))
#1.63 ms ± 526 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

_ = j_update_tree_new(trees_Ns,stack_p(indx),priority_dmy)  # compile
%timeit jax.block_until_ready(j_update_tree(trees_Ns,stack_p(indx),priority_dmy))
#462 µs ± 76.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

_ = v_update_tree_new(trees_Ns,stack_p(indx),priority_dmy)  # compile
%timeit jax.block_until_ready(v_update_tree(trees_Ns,stack_p(indx),priority_dmy))
#3.37 ms ± 98.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

As you can see, there is no distinctive difference of using lax.scan compared to the original for-loop based function. What might be the reason and how should I fix it?

[jakevdp]

I generally wouldn't expect JIT-compiled scan to be faster than a JIT-compiled Python for loop at runtime. But you'll find that as the size of the loop grows, the for loop solution will take much longer to compile—roughly scaling as the square of the number of iterations—due to the fact that JAX tracing flattens all Python control flow. This compile time cost is the main reason to prefer scan over for loops in practice.