Torch2Jax vs Torchax?

[adam-hartshorne]

I have just come across the new Torchax library from Google for cross JAX / Pytorch interactions. I was just wondering what the pros / cons might between the two libaries for using it to call PyTorch code from within a JITed Jax model?

[rdyro]

The answer here: google/torchax#55 is fully correct.

torch2jax is a way to wrap torch computations in a black box which jax doesn't see into, but can execute efficiently.

Despite the black box approach torch2jax defines a backwards pass, so you can differentiate through the whole program and should run multi-device comptuations correctly.