Provide output buffer for pure_callback result about jax

[superbobry]

There is a preallocated buffer on the XLA side, but it is not currently passed to the pure_callback.

@hawkinsp should we have an API for this, wdyt?

[superbobry]

Hey @Joshuaalbert, sorry for the silence. Here is a quick update

• @yashk2810 is working on a proposal to allow "write once" mutation for jax.Arrays which are currently immutable (and thus cannot be used for output buffers);
• I'm exploring a parallel idea -- dlpack.callback, a new callback API using DLPack capsules for both inputs and outputs.

[superbobry]

Another quick update: I prototyped dlpack.callback, but after discussing it with a few JAX team members, I decided not to move forward with it as JAX has too many callback APIs already.

Instead, the plan is to change existing callback APIs to support mutable_results=. I am waiting on a few changes in XLA FFI, but once they land, it should be fairly straightforward to implement this.

[superbobry]

No news yet in a sense that none of my attempts landed. Hopefully, in the coming weeks :)

[dfm]

@Joshuaalbert — JAX doesn't currently offer a public API for this (customizing multiple transforms for a single callable), but it's on our radar. I'd say that this is off topic for this issue thread, but feel free to open another with more info about your use cases for using both JVP and VJP for one callback!

[Joshuaalbert]

Also, is it safe to reuse the buffer of input values, e.g. like

def _fn(x):
    np.square(x, out=x)
    return x

result_shape_dtype = jax.ShapeDtypeStruct(
    shape=np.shape(x),
    dtype=x.dtype
)

return jax.pure_callback(_fn, result_shape_dtype, x, vectorized=True)

EDIT: That would be a nope as it gives jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: ValueError: output array is read-only

[Joshuaalbert]

Also, the main reason this would be helpful is because the output buffers for the science applications I'm working on are really large so if there is already one allocated by XLA it would save lots of memory to use that one.

[Joshuaalbert]

@superbobry @hawkinsp any update available for this?

[Joshuaalbert]

Hi @superbobry, any news?

[Joshuaalbert]

@superbobry can you check out JAXbind which offers a way to specify both JVP and VJP for external callbacks? I argue that it should be also possible within JAX to specify both.