Chris Cundy

`host_callback.call` fails on multi-gpu machine

4

If I run the following code: ``` from jax.experimental import host_callback import numpy as np from jax import pmap, jit, partial, ShapeDtypeStruct def host_fn(x): return x x = np.ones(4, dtype=np.float32)...