Torchax vs Torch2Jax ?

[adam-hartshorne]

Does Torchax have an advantage over using Torch2Jax (https://github.com/rdyro/torch2jax) if we are wanting to call pytorch functionality from within a JITTed JAX program?

Will the way Torchax works cause issues if one wanted to try to pass the Pytorch tensors to Pykeops from within a JittableModule?

https://www.kernel-operations.io/keops/python/LazyTensor.html

[qihqi]

Hi @adam-hartshorne,

I took a brief pass on reading how torch2jax works; and it looks like it's doing the following:

1. save the torch callable reference as-is, in somewhere the C++ side can access
2. Pass it to C++ side through FFI, and use libtorch to run the torch callable
3. Enough registration / partition params etc to route stuff.

cc. @rdyro please correct me if the above understanding is not correct.

Contrast to torchax, which is to implement each torch Aten op in JAX directly, and use tensor subclass to route: https://google.github.io/torchax/user_guide/how-it-works/

In summary, the torch math in torch2jax is executed by libtorch itself which have the following pros & cons:

Pros:

• No numerics issue: whatever math done in torch eager mode is the same code running under jax.jit
• No op coverage issue: all ops works by definition (unless it also dont work in torch)

Cons:

• Doesn't run on devices that libtorch itself doesnt work, but Jax works. Namely Google TPU. I also tried it torch2jax on Mac and was not successful (due to need to build extension).
• Is not a viable export path to StableHLO: if you do jax.jit(f).lower(...).as_text() with a torch2jax wrapped computation, you'd likely get a graph with a single FFI call to a Cpp function that calls libtorch (disclaimer: I havent tried myself, but this profile seems to imply so: https://github.com/rdyro/torch2jax/blob/main/images/data_parallel.png). If you do the above in torchax, you'd get a StableHLO graph.
• Less tight integration with Jax: the ffi node is a black box to Jax. For example, the above profiler graph will not see the individual op that composed on wrapped torch model.

On LazyTensor, It should work as long as the Jax-jit wrapped function materialize the Lazy tensor, and return materialized plain tensor (which case would be the torchax.tensor.Tensor subclass), so jax can trace everything it need to trace.
However, using Lazy Tensor is unlikely to add any performance improvements because jax-jit also does full graph compilation. So I'd advice to stick with plain Pytorch.

[rdyro]

Hey @qihqi that's exactly right!

[adam-hartshorne]

Thanks for your detailed reply.