Wrap 3rd party memory as jax.numpy.array?

[bkkgbkjb]

Hi all

I wonder if I could wrap an external, 3rd party memory chunks as jax.numpy.array such that it could enjoy the ecosystem of jax?

For example, certain simulation platform like Isaacgym performs their own simulation logics on CPU/GPU. Then port it to/from pytorch tensor for front-end users. example

I wonder if we could do the simliar thing for JAX? I know dlpack could be employed in some ways but wonder if we have more direct way to reduce conversion overhead as much as possible?

BTW, I don't need to differentiate through this array. So it's purely an interface between front-end dl library and back-end simulation platform.

[jakevdp]

The main mechanism JAX provides for this is via dlpack. See the jax.dlpack module, and also the NumPy-style API at jax.numpy.from_dlpack.

[bkkgbkjb]

Thanks for the reply!

This is just what I need

But with some quick lookup, it seems the latest jax (version 0.5.2) doesn't support DLPack v1.x?

If so, which version of dlpack should I study and use?

big thanks in advance

[pearu]

But with some quick lookup, it seems the latest jax (version 0.5.2) doesn't support DLPack v1.x?

This is not true. DLPack v1.1 was released just 2 days ago and there is openxla/xla#23653 on the line. JAX/XLA supports DLManagedTensorVersioned that was introduced in DLPack v1.0.

What resources did you use to conclude that jax does not support DLPack v1.x?

[bkkgbkjb]

wow, that is awesome

my premature conclusion is drawn from my personal understanding of the note presented in the docs

Specifically,

While JAX arrays are always immutable, DLPackManagedTensor buffers cannot be marked as immutable, and it is possible for processes external to JAX to mutate them in-place. If a DLPack buffer derived from a JAX array is mutated, it may lead to undefined behavior when using the associated JAX array. When JAX eventually supports DLManagedTensorVersioned (DLPack 1.0), it will be possible to specify that a buffer is read-only.

Maybe this is a outdated note ?

[pearu]

to_dlpack is a legacy DLPack interface and its usage should be deprecated.

For wrapping a 3rd party dlpack-compatible array object obj as JAX array, use jax.numpy.from_dlpack(obj).

For wrapping a JAX array arr as a 3rd party array object, use <thirdparty-library>.from_dlpack(arr).

[bkkgbkjb]

ooook, so much thanks!

really appreciate your suggestions and reply!