Need guidance on exposing mutable external array to JAX in-place on CUDA platform

[jjyyxx]

My current understanding is that:

• the recommended way to import existing CUDA memory to JAX is either exposing __cuda_array_interface__ or __dlpack__()+__dlpack_device__() (works in-place if used properly);
• The way to opt-in mutable access is calling jax.array_ref on an JAX array to obtain a mutable ArrayRef;
• Typed XLA FFI is the modern way to bridge and extend JAX with custom functionality.

A clarification first: Some external systems, such as google-deepmind/mujoco_warp, instead of importing external memory, let JAX allocate and manage arrays for tighter and smoother integration. But in my use case, I cannot go this way, since memory allocated by JAX has unmet property for the external system I'm trying to interface with. Another solution is to maintain a copy at JAX side, but this leads to unnecessary memory consumption and copy overhead if the array is large.

Here I have several confusions:

1. jax.array_ref copies the jax.Array internally (understandable to maintain the fundamental immutable semantics), preventing me from obtaining a mutable reference to imported array in-place. A hack is to call from jax._src.state.types import AbstractRef; ref = jax.ArrayRef(AbstractRef(arr.aval), arr), and it does work, but using internal API is what I'd like to avoid. So here are two questions:
   • Is it possible to provide such an API (less safe than the current copy variant)?
   • And what are semantics of mutability in JAX (ordering with other statements? visibility if more than one ArrayRef is aliasing under the hood?)
2. ArrayRef cannot be passed to FFI call directly (some confusing state discharge error). In this case, how to pass the exact array to FFI function? jax.ref.freeze leads to another copy, while passing a jax.Array has no guarantee that the device pointer received by FFI is identical to the intended one. Since the array is external to JAX, I can, of course, pass it as context outside JAX, but I worry about uncaptured implicit dependency leading to ordering issue.

[jakevdp]

array_ref is suitable for JAX-internal mutable array references. I don't believe it is designed to support mutable views of buffers that originate outside of JAX, and I don't believe usage with FFI is one of its design goals (@mattjj or @dougalm may be able to comment more).

[mattjj]

Actually, we do want ArrayRefs to play well with the FFI and to model buffer ownership at the runtime level. But we haven't actually gotten there yet!

@hawkinsp @pschuh is there anything useful to say at this point?