Have any of you seen benchmarks comparing the performance of `@jax.jit` and `@torch.compile`, especially when using functional PyTorch code? Are the performance differences big? Small? Do they depend a lot on what you’re doing?
Have any of you seen benchmarks comparing the performance of `@jax.jit` and `@torch.compile`, especially when using functional PyTorch code? Are the performance differences big? Small? Do they depend a lot on what you’re doing?
I have not gone into the internals, but from the documentation they seem to be behaving in a very similar manner in the sense that they both do symbolic execution by sending a tracer object, recording what’s happening, and then emitting compiled code.
JAX specifically takes into consideration shapes, and it seems that torch does it as well. Both might do more jitting if it’s necessary, eg different input shapes.
I have the hunch that Jax should be faster because it relies on XLA primitives, eg xla-backed for-loops and scans, and I am not sure how PT handles that given it most likely tries to remain backwards compatible and rely on python semantics over some compiler primitives.