News for hackers. Stuff that matters.

A comment on Supercharged high-resolution ocean simulation with Jax.

Re:Supercharged high-resolution ocean simulation with Jax by bertr4nd on Sunday December 05, 2021 @07:54PM (#29454972)
In a shameless plug, I want to note that running these sorts of workloads on CPU using Pytorch got much faster (some results on a benchmark from this post’s author’s suite in [0]) in the most recent torch release thanks to the addition of a JIT compiler. Obviously there’s much to recommend Jax (the XLA compiler is quite excellent), but it’s nice to have some choice in the space.

[0] https://www.linkedin.com/feed/update/activity:68640106214579...

[ Parent | on HN ]
Re:Supercharged high-resolution ocean simulation with Jax by dionhaefner on Monday December 06, 2021 @03:19AM (#29457444)
True, but unfortunately Pytorch is not quite there yet when it comes to more complex benchmarks:

https://github.com/dionhaefner/pyhpc-benchmarks#example-resu...

JAX really is the only library that comes close to low-level code on CPU, almost always (that I've tried).

[ Parent | on HN ]
Re:Supercharged high-resolution ocean simulation with Jax by bertr4nd on Monday December 06, 2021 @06:21AM (#29458362)
Interesting, I thought pytorch was a bit more competitive on those other benchmarks (but admittedly it’s been a while since I looked). Slicing shouldn’t be a fundamental problem, but perhaps there are some important details that have been overlooked. Thanks for pointing it out!
[ Parent | on HN ]