Parallelization of pure_callabck

[qres]

Is it possible to execute a pure_callback in parallel? The following executes all calls serially when executed with env XLA_FLAGS="--xla_force_host_platform_device_count=2 --xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" OPENBLAS_NUM_THREADS=1.

from jax.config import config
config.update("jax_platform_name", 'cpu')
import jax
import jax.numpy as jnp
from jax.sharding import PositionalSharding
import numpy as np

key = jax.random.PRNGKey(1)
data = jax.random.uniform(key, (10, 1000)) # we want to parallelize over the first dimension

def callback(ix, data):
    jax.debug.print('{}', ix)

    # some scipy code using jit-ed jax functions

    # dummy code
    import time
    time.sleep(1)
    result = np.sum(2*data)
    success = np.array(True)

    return success, result

result_shape = [
    jax.ShapeDtypeStruct((), np.dtype('bool')),
    jax.ShapeDtypeStruct((), np.dtype('float32')),
]
p_callback = lambda ix, data: jax.pure_callback(callback, result_shape, ix, data)

sharding = PositionalSharding(jax.devices())

result = jax.vmap(p_callback)(
    jax.device_put(jnp.arange(data.shape[0]), sharding),
    jax.device_put(data, sharding.reshape(-1, 1)),
)
result[0].block_until_ready()

print()
print(result)

(Using jit around the vmap gives me the error Command terminated by signal 11 (jax version 0.4.12 on Ubuntu 20.04))

The pure_callback is rather expensive and should therefore be parallelized over multiple cores for different parameters.
I'm also open for other ideas how this can be parallelized. I tried python's multiprocessing.Pool.map but it complained that jit-ed jax-functions could not be pickled.

Thank you for your help!

[jakevdp]

A pure_callback is a function that executes on the host at runtime, and host execution occurs in a single CPU process. So in that sense, no, you cannot execute a pure_callback in parallel.

On the other hand, your host-side function is just a Python function, which could call out to any parallel computation framework that the host Python process has access to (including another JAX process, python multiprocessing, pytorch, dask, etc.), so in that sense you could execute a pure_callback in parallel, but it has nothing to do with pure_callback per se.

Does that make sense?