JAX's external callbacks

[Joshuaalbert]

This package is very similar to JAX's provided, and improving, external callback functionality.

I also urge you to show support for jax-ml/jax#20701 to get JAX provide output buffers.

[mreineck]

If I understand the discussion in #10 correcty, JAXbind's functionality goes beyond what JAX's external callbacks provide.

Being able to specify your own output buffers is certainly an attractive feature; if it can be made to work with JAXbind, I hope we can support this soon. For the moment, the focus is on finishing the code paper, however.

[Edenhofer]

JAXbind already passes on the pre-allocated buffer from XLA to the user. IMO this is a non-issue here. Please see the paper for a discussion of the differences between JAXbind and JAX's external callback API (most importantly the support for customizing both JVP and VJP in JAXbind).

I agree that jax-ml/jax#20701 would be neat to have.

[Joshuaalbert]

JAX's external callbacks also allow specifying a custom jvp and vjp using jax.custom_jvp/jax.custom_vjp. See here for example specifying a jvp.

[Joshuaalbert]

Ah, I see, so it's related to allowing both JVP and VJP #10 (comment). Nice. My only comment is then that it would be better to follow the API of JAX's external callbacks, as much as possible since there is then a better chance they'll incorperate it into JAX's source, which I would argue should be the goal of JAXbind. Otherwise, there is just competing offers and redundant efforts.

[Edenhofer]

We reached out to the JAX developers regarding this here jax-ml/jax#19179 .

I am not sure whether a unified interface is easily possible given that JAXbind needs to combine both pure_callback, custom_jvp and custom_vjp in one. I can imagine a future where JAXbind will simply be superseded by tools in JAX itself (e.g. if custom transpositions would be added and issues like the one you raised are solved). However, we are not there yet otherwise we wouldn't have felt the need to code JAXbind.

[SepandKashani]

@Joshuaalbert You mentioned wanting to try out JAX's jax.experimental.buffer_callback() function. If so, what is your take on it?

I have only been aware of it since last week (since it is not exposed in the HTML docs...), but it does allow passing DLPack-compatible arrays to JAX without forcing a host synchronization. What it lacks is support for jvp/vjp, but I assume support for these transforms can be added to the primitive similar to how JAXbind does it.

I will continue playing around with it to assess how difficult this task would be. If conclusive, perhaps this would allow to simplify JAXbind's internals?

[Edenhofer]

From my understanding buffer_callback will not significantly simplify the code in JAXbind as we still need our own custom primite to define both JVP and VJP.

[mreineck]

Is this issue still active or should we close it?