Is any JAX external callback async/non-blocking?

[susuzheng]

We want to have a callback inside jitted code to run in parallel with other TPU computations. To illustrate, take the following as an example. We hope the first and second line could run in parallel, one on CPU and one TPU. Then, run the third line when both are ready. Would this be possible with a single device setup?

# inside a jitted function running on TPUs.
mask = jax.pure_callback(...)
logits = mdoel()
logits.at[~mask].set(-1)

[jakevdp]

Pinging @sharadmv who may know more: I believe that callbacks are non-blocking on TPU, but ultimately the details will be up to the XLA scheduler.