jax

0eace80a - Fix experimental host callback on multi-host (#3200)

Fix experimental host callback on multi-host (#3200) * Fix experimental host callback on multi-host Hosts can only access the outfeed queue for local devices, while `api.devices` returns all devices in the system. * Update host_callback.py