3 ms·
Because there's a bunch of layers between the high level code and the output? For simple stuff, we can compare JAX to PyTorch on a 4090, and JAX seems faster b
by VHRanger 4y ago
Because there's a bunch of layers between the high level code and the output?
For simple stuff, we can compare JAX to PyTorch on a 4090, and JAX seems faster by 10-50%. It's way way way faster on TPU.
That said, JAX is apparently a pain to work with in comparison to PyTorch (disclaimer I haven't used JAX enough to opine yet).
I'd also want a comparison between JAX and Taichi-lang for parallel stuff, even though the problem domains aren't exactly aligned.
- chillee 4y ago> For simple stuff, we can compare JAX to PyTorch on a 4090, and JAX seems faster by 10-50%. It's way way way faster on TPU. What benchmarks are you looking at here?
- VHRanger 4y agoIndividual user benchmarks on github and 2-3 blogs. Not super scientific, Im sorry to say, but methodologically benchmarking this is really hard.