JAX - PyTorch interoperability

[Toni-SM]

Hi.

I am getting the following deprecation warning when doing JAX - PyTorch interoperability via dlpack:

DeprecationWarning: Calling from_dlpack with a DLPack tensor is deprecated. The argument to from_dlpack should be an array from another framework that implements the __dlpack__ protocol.

import jax.dlpack as jax_dlpack
import torch.utils.dlpack as torch_dlpack
# ...
jax_dlpack.from_dlpack(torch_dlpack.to_dlpack(tensor.contiguous()))

Question: What, if any, is the "new" way of doing JAX - PyTorch interoperability?

[adam-hartshorne]

These two packages might help,

https://github.com/rdyro/torch2jax
https://github.com/samuela/torch2jax

I have personally used the first one a lot, it is now using the new JAX FFI interface and found it extremely efficient / the maintainer is very helpful.

Going the other way,

https://github.com/lucidrains/jax2torch

[Toni-SM]

Hi @adam-hartshorne

Thank you for those indications.
For now, I am going to use @hawkinsp solution to avoid installing extra dependencies, but I will take a look at these modules

[hawkinsp]

IIRC, it's jax_dlpack.from_dlpack(tensor.contiguous()) these days.

[hawkinsp]

In [2]: import jax
In [3]: import torch
In [4]: x = torch.arange(5)
In [5]: jax.dlpack.from_dlpack(x)
Out[5]: Array([0, 1, 2, 3, 4], dtype=int32)