Supercharged high-resolution ocean simulation with Jax
dionhaefner.github.io
Supercharged high-resolution ocean simulation with Jax
1–10 of 49 posts
Re: Supercharged high-resolution ocean simulation with Jax
#2Awesome! One question immediately comes to mind. Any interest in doing this stuff with Julia? You're basically the epitome of their target audience: a scientific computing expert who does HPC with differentiable programs.
Re: Supercharged high-resolution ocean simulation with Jax
#3Awesome! One question immediately comes to mind. Any interest in doing this stuff with Julia? You're basically the epitome of their target audience: a scientific computing expert who does HPC with differentiable programs.
There's an earlier blog post by the same author where they discuss three possible ways of moving away from the Fortran/C status quo towards higher-level models. They mention Julia as one of the routes, but not the one they decided to take: https://dionhaefner.github.io/2021/04/higher-level-geophysic...
Re: Supercharged high-resolution ocean simulation with Jax
#4In 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...
Re: Supercharged high-resolution ocean simulation with Jax
#5What made you choose JAX over Julia? I'm interested in this question, because I have been thinking about transitioning to Julia but have always hesitated to make the move, since overall the Python ecosystem still seems way ahead in terms of visualization and toolchain.
Also, would you expect JAX acceleration to work well with other types of discretization, such as spectral methods?
Re: Supercharged high-resolution ocean simulation with Jax
#6My question is how much of the operations in JAX here can be done with reduced precision and can utilize training accelerators i.e. TPUs. I've noticed a lot of research coming out in physics, where everything is simulated in at least double float, being augmented with ML approaches where precision is traded for dynamic range.
Re: Supercharged high-resolution ocean simulation with Jax
#7What made you choose JAX over Julia? I'm interested in this question, because I have been thinking about transitioning to Julia but have always hesitated to make the move, since overall the Python ecosystem still seems way ahead in terms of visualization and toolchain. Also, would you expect JAX acceleration to work well with other types of discretization, such as spectral methods?
From the post:
> JAX on GPU outperforms everything
I've only skimmed through the blog post, but it feels that the GPU acceleration without a need to write any custom code, was the primary reason to choose JAX.
Re: Supercharged high-resolution ocean simulation with Jax
#8What made you choose JAX over Julia? I'm interested in this question, because I have been thinking about transitioning to Julia but have always hesitated to make the move, since overall the Python ecosystem still seems way ahead in terms of visualization and toolchain. Also, would you expect JAX acceleration to work well with other types of discretization, such as spectral methods?
If numpy is a good fit for it, JAX is a good fit for accelerating it, basically. I think of it as numpy plus program transformations, such as differentiation, JIT, parallelization, compiling to XLA, compiling to TPU/GPU, etc.
The magic of JAX is that it keeps all that stuff about as simple as writing numpy code.
Re: Supercharged high-resolution ocean simulation with Jax
#9Awesome! One question immediately comes to mind. Any interest in doing this stuff with Julia? You're basically the epitome of their target audience: a scientific computing expert who does HPC with differentiable programs.
There's an earlier blog post by the same author where they discuss three possible ways of moving away from the Fortran/C status quo towards higher-level models. They mention Julia as one of the routes, but not the one they decided to take: https://dionhaefner.github.io/2021/04/higher-level-geophysic...
"On the other hand, Julia’s focus on scientific applications is both blessing and curse. In this day and age, a lot of the progress in computing is driven by applications outside academia (mostly through machine learning)." This seems like a crazy mis-read to me. Julia is probably the language that has the best integration of differential equations and machine learning. Jax closes the gap a little, but is still way behind.
For example https://gist.github.com/ChrisRackauckas/62a063f23cccf3a55a4a... shows a pretty simple case where DifferentialEquations.JL is 6x faster at gradient calculations than Jax.
Re: Supercharged high-resolution ocean simulation with Jax
#10like some other commenters here, https://github.com/CliMA/Oceananigans.jl immediately comes to mind, maybe it would be fun to compare projects on this scale between JAX/Julia.
> JAX offers more than just a JIT compiler: JAX functions are also differentiable
if the downstream library is completely implemented in JAX (numba) ecosystem. Similar for Julia, except implementing fast code in Julia is natural, doesn't involve debugging 3 compilers (Cpython, Numba, Jax). Many python library is only differentiable because the 100x more effort were put in writing C/C++ backend, binding to python, and writing chain rules for foreign functions.
I would imagine Julia to be a good fit for this direction in the future!