Ever Wanted to Run Torch Code From Jax - New Torch2Jax Library

[adam-hartshorne]

This isn't my work, but I think it deserves a much wider viewing as it is an absolutely fantastic bit of work and incredibly useful.

This package is designed to facilitate no-copy PyTorch calling from JAX under both eager execution and JIT. It leverages the JAX C++ extension interface, enabling operations on both CPU and GPU platforms. Moreover, it allows for executing arbitrary PyTorch code from JAX under eager execution and JIT.
https://github.com/rdyro/torch2jax