Ever Wanted to Run Torch Code From Jax - New Torch2Jax Library #16817
Unanswered
adam-hartshorne
asked this question in
General
Replies: 0 comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
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
Beta Was this translation helpful? Give feedback.
All reactions