Parse warp array to google jax

Hello intercommunication between PyTorch tensor and warp arrays is present, is there any way to parse jax devicearray into and from warp arrays?

Hi @jakub.mitura14. Checking with the devs.

1 Like

Hi Jakub,

Currently we don’t offer JAX interop out of the box - we do support the _cuda_array_interface_ protocol that is quite commonly used, but JAX requires dlpack.

We will look at adding dlpack support in the near future.

Best,
Miles

1 Like

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.