jax

cf6655fd - [JAX] Add `py_pjrt_host_callback` extension to jaxlib

[JAX] Add `py_pjrt_host_callback` extension to jaxlib Exposes `create_hlo_output_callback` to Python to allow creating `HloOutputCallback`s from Python functions in jaxlib. PiperOrigin-RevId: 934031057