Optimizing XLA Compilation when using a tuple of functions as an argument in a scan function

[christophedessers]

Hello,

I’m encountering a significant increase in XLA compilation time when using a tuple of functions as an argument in a scan function which uses a switch to select some function to apply.

Specifically, the issue arises when I scale the system by duplicating these functions. If I understand correctly, creating N functions would lead JAX to compile N functions. However, in my case, I’m only working with two distinct functions (fct1, fct2) which are partial and both of them only have 1 value as a "curried" parameter. Therefore, although I replicate these functions many times, I hoped that JAX would see there are actually only two 2 unique functions, not N.
Here is the simplified code :

import time
from typing import Callable
import jax.numpy as jnp
from jax import jit
from jax.tree_util import Partial
from jax.lax import scan, switch

# Functions' library

def fct1(u:jnp.ndarray, u_prev:jnp.ndarray, params:tuple):
    return u + u_prev + params[0]

def fct2(u:jnp.ndarray, u_prev:jnp.ndarray, params:tuple):
    return u**params[0] + u_prev

def apply_one_fct(index:int, U:jnp.ndarray, U_prev:jnp.ndarray,
                  params:tuple,
                  fct_list:tuple, i_fct:jnp.ndarray):
    index_fct = i_fct.at[index].get()
    # Note the presence of a switch. We are basically using "i_fct" to
    # select which function (found in "fct_list") to apply.
    return switch(index_fct, fct_list, U, U_prev, params)

def run(U_prev:jnp.ndarray, U:jnp.ndarray,
        fct_select:Callable[[int, jnp.ndarray, jnp.ndarray, tuple], tuple],
        params:tuple):
    nb_unknowns = jnp.size(U)
    out = jnp.zeros(nb_unknowns)
    for i in range(len(params)):
        out = out.at[i].set(fct_select(i, U.at[i].get(), U_prev.at[i].get(), params[i]))
    return out, out

# Test
print("Initialisation started...")

# Model parameters
param_fct1 = 0.05
param_fct2 = 0.01

# Simulation parameters
nb_replica = 5_000
nb_time_steps = 1_000

# Initial conditions
u_0 = jnp.linspace(0, 1, nb_replica*2)
u = u_0
u_tot = jnp.zeros((nb_time_steps, nb_replica*2))

# Build all the tuples to prepare the simulation
fcts = (fct1, fct2)
params_model = nb_replica * ((param_fct1,), (param_fct2,))
indices_fct = jnp.array([0, 1]*nb_replica)

# Compilation of the functions
fct_switch = jit(Partial(apply_one_fct, fct_list=fcts, i_fct=indices_fct))
solve_fct = jit(Partial(run, fct_select=fct_switch, params=params_model))


# Simulation
print("Running the simulation...")
begin_t = time.time()

# The following line boils down to:
#    jax.scan(
#       solve_fct(
#          run(
#            fct_switch(
#               jax.switch(list_of_functions, list_of_indices, list_of_parameters)))))

final, result = scan(solve_fct, u_0, xs=u_tot)
print("final = ", final.block_until_ready())
print(f"The time of the program is : {time.time() - begin_t} s.")

Here’s when the problem arises: As the number of functions duplicates increases, the XLA compilation time for the scan function grows rapidly (for example, with 5000 replicas, compilation takes about 2 minutes). I suspect JAX doesn't understand that there are 2 unique functions.

My questions:

1. Is there a more efficient way to help JAX optimize the compilation process, especially when only two or a small set of functions are being reused?
2. Why does XLA compilation time increase so dramatically with the number of function duplicates, while MLIR compilation time seems to remain relatively stable?

Here are the results of the “xla_dump_to” flag containing the pre-optimization HLO files :
dump_Simple_test.zip

I’m new to JAX, so I might be missing some key optimization strategies. Any insights or suggestions would be greatly appreciated!

Thank you for your time.