Running multiple functions in parallel in JAX

[tomas-teijeiro]

I have a bunch of JAX-jitted functions that I want to run in parallel with the same input (a JAX array). However, I couldn't find any proper way of doing this, neither within the JAX framework or using standard Python multiprocessing APIs. This is a minimum working example of what I want:

import jax
import jax.numpy as jnp
import numpy as np

#Array with many different functions (in reality they are all different)
functions = [jnp.sum]*20 + [jnp.prod]*20 + [jnp.mean]*20 + [jnp.median]*20

@jax.jit
def run_multiple_functions(x):
    return jnp.vstack([f(x) for f in functions])


x = np.random.rand(10000)

run_multiple_functions(x)

Everything works well in this example, but the execution time of run_multiple_functions is the sum of the execution time of all of them. If I try to use a ProcessPoolExecutor inside run_multiple_functions, the code breaks due to a pickle problem, and if I use a ThreadPoolExecutor I don't get any benefit due to the GIL.

Any ideas on how I could do this? Many thanks in advance!!

[ASEM000]

I am not sure if this is optimal for your case, but you can try something like shard_map:

import os
import jax
import jax.numpy as jnp
import functools as ft
from jax.sharding import PartitionSpec as P
from jax.experimental.shard_map import shard_map
import numpy.testing as npt
mesh = jax.make_mesh([2], ["P"])


def f1(x, y):
    return x @ y

def f2(x, y):
    return -x @ y


@ft.partial(
    shard_map,
    mesh=mesh,
    in_specs=(P(), P()),
    out_specs=P("P"),
)
def sharded_f(x, y):
    index = jax.lax.axis_index("P")
    return jnp.where(index == 0, f1(x, y), f2(x, y))


def unsharded_f(x, y):
    return jnp.vstack([f1(x, y), f2(x, y)])


D = 2048
x = jnp.ones((D, D))
y = jnp.ones((D, D))

npt.assert_allclose(sharded_f(x, y), unsharded_f(x, y), atol=1e-4)