Pipelined calculation involving scan and pure_callback

[mfschubert]

This question is in regards to pipelined calculations involving jax.lax.scan and jax.pure_callback. I have an expensive calculation with two parts:

• one part is non-jax, runs on CPU, and is called via pure_callback
• the second part is jax and runs on GPU

Both calculations are used in a scan operation, as shown in the example below.

import jax
import jax.numpy as jnp
import time

@jax.jit
def slow_callback_fn(key):
    # Dummy calculation that runs on CPU and takes 500 ms.
    def _fn(x):
        time.sleep(0.5)
        return x
    
    return jax.pure_callback(_fn, key, key)

@jax.jit
def dummy_jax_fn(key):
    # Dummy calculation that runs on GPU and takes roughly 500 ms.
    m = jax.random.normal(key, (5000, 5000))
    return jnp.mean(jnp.linalg.inv(m @ m @ m))

def serial_scan(carry, key):
    del carry
    return (), dummy_jax_fn(slow_callback_fn(key))

keys = jax.random.split(jax.random.PRNGKey(0), num=10)

t0 = time.time()
_, xs = jax.lax.scan(serial_scan, (), keys, unroll=1)
xs.block_until_ready()
print(time.time() - t0)

I am hoping to pipeline the calculation to speed things up, as follows:

def pipelined_scan(carry, key):
    return slow_callback_fn(key), dummy_jax_fn(carry)

t0 = time.time()
init = slow_callback_fn(keys[0])
carry, xs_pipelined = jax.lax.scan(pipelined_scan, init, keys[1:], unroll=1)
xs_pipelined = jnp.concatenate([xs_pipelined, jnp.asarray([dummy_jax_fn(carry)])])
xs_pipelined.block_until_ready()
print(time.time() - t0)

I would expect the compute time to be cut in half here, since the slow_callback_fn(key) and dummy_jax_fn(carry) take an equal amount of time and can run independently. However, this doesn't seem to be the case in practice.

Is this expected? Is there some other way I can force these two calculations to run in parallel?