Determine how to support JAX's callback system

Open rsuderman opened this issue 2 years ago • 3 comments

JAX supports registering callbacks that can be triggered via the backend environment. This may require implementation via the mhlo.custom_call system or an alternative integration method via the PJRT events system.

Feb 27 '23 23:02 rsuderman

Or via something like https://github.com/openxla/iree/tree/main/samples/py_custom_module if this is callback to Python.

Feb 28 '23 16:02 jpienaar

Is it this:https://jax.readthedocs.io/en/latest/jax.experimental.host_callback.html ?

Mar 22 '23 00:03 vmurali

Can this issue be closed

May 14 '24 03:05 vmurali