Hacker News Board Hacker News, browsed like a classic message board
Live data from Hacker News

Supercharged high-resolution ocean simulation with Jax

dionhaefner.github.io

Post reply on HN
Print view View original

11–20 of 49 posts

Re: Supercharged high-resolution ocean simulation with Jax

#11

Post by moelf »

krasin wrote: post #7
geoalchimista wrote: post #5

What 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.

>without a need to write any custom code, was the primary reason to choose JAX.

this is even more "free" in Julia, JAX at least need to worry when foreign call happens (library not derived from Numpy/JAX ecosystem, or outright C/C++ binding without JAX rules).

Re: Supercharged high-resolution ocean simulation with Jax

#12

Post by adgjlsfhk1 »

oakridge wrote: post #6

My 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.

I would recommend checking out https://www.youtube.com/watch?v=GiSsoA1udUk. It shows that you can can do climate models with 16 bit numbers.

Re: Supercharged high-resolution ocean simulation with Jax

#13

Post by krasin »

moelf wrote: post #11
krasin wrote: post #7

Earlier quoted context omitted.

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.

>without a need to write any custom code, was the primary reason to choose JAX. this is even more "free" in Julia, JAX at least need to worry when foreign call happens (library not derived from Numpy/JAX ecosystem, or outright C/C++ binding without JAX rules).

For someone like me who is familiar with JAX, but only recently starting to consider to pay attention to Julia (it does have a momentum), would it be possible to provide any good examples of using GPU / multiple GPUs from Julia?

I've tried to search on my own, but only a way to write CUDA-dependent code: https://juliagpu.gitlab.io/CUDA.jl/usage/multigpu/

Re: Supercharged high-resolution ocean simulation with Jax

#14

Post by mjn »

adgjlsfhk1 wrote: post #9
mjn wrote: post #3

Earlier quoted context omitted.

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 b…

I read that as being about what language industry uses to write ML applications, not about technical feasibility of integrating machine learning methods into a codebase. Put differently: industry most often uses Python (especially in ML), therefore the author wants to target Python in order to maximize uptake outside of academia. They even admit that doing it in Python is technically harder than doing it in Julia ("Unfortunately, this type [Type III] is also the hardest to get right"), but consider it worth the trouble for the broader accessibility.

(That's more or less the direction I've been going with research code lately too, so I can sympathize, although I'm not entirely happy with the situation and definitely also sympathize with the Julia folks being unhappy about it.)

Re: Supercharged high-resolution ocean simulation with Jax

#15

Post by moelf »

krasin wrote: post #13
moelf wrote: post #11

Earlier quoted context omitted.

>without a need to write any custom code, was the primary reason to choose JAX. this is even more "free" in Julia, JAX at least need to worry when foreign call happens (library not derived from Numpy/JAX ecosystem, or outright C/C++ binding without JAX rules).

For someone like me who is familiar with JAX, but only recently starting to consider to pay attention to Julia (it does have a momentum), would it be possible to provide any good examples of using GPU / multiple GPUs from Julia? I've tried to search on my own, but only a way to write CUDA-dependent code: https://juliagpu.gitlab.io/CUDA.jl/usage/multigpu/

I'm not familiar with multi-GPU setup in general. GPU programming in Julia has the advantage that naive operation doesn't even need to be GPU-aware (for writers), since GPU arrays (of any vendor backend) conforms the AbstractArray interface.

If you're advanced library writer, you can leverage: https://juliagpu.github.io/KernelAbstractions.jl/stable/#Wri... which allows you to write kernel, in Julia, that compiles efficiently with rest of native Julia code, and that works cross-vendor!

Back to multi GPU, it seems there's: https://clima.github.io/OceananigansDocumentation/stable/app... which is MPI based?

Re: Supercharged high-resolution ocean simulation with Jax

#16

Post by adgjlsfhk1 »

krasin wrote: post #13
moelf wrote: post #11

Earlier quoted context omitted.

>without a need to write any custom code, was the primary reason to choose JAX. this is even more "free" in Julia, JAX at least need to worry when foreign call happens (library not derived from Numpy/JAX ecosystem, or outright C/C++ binding without JAX rules).

For someone like me who is familiar with JAX, but only recently starting to consider to pay attention to Julia (it does have a momentum), would it be possible to provide any good examples of using GPU / multiple GPUs from Julia? I've tried to search on my own, but only a way to write CUDA-dependent code: https://juliagpu.gitlab.io/CUDA.jl/usage/multigpu/

https://www.juliapackages.com/p/pencilarrays is a really good tool to do this type of stuff automatically for some applications.

Re: Supercharged high-resolution ocean simulation with Jax

#17

Post by moelf »

mjn wrote: post #14
adgjlsfhk1 wrote: post #9

Earlier quoted context omitted.

"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 b…

I read that as being about what language industry uses to write ML applications, not about technical feasibility of integrating machine learning methods into a codebase. Put differently: industry most often uses Python (especially in ML), therefore the author wants to target Python in order to maximize uptake outside of academia. They even admit that doing it in Python is technically harder than doing it in Julia ("U…

hopefully if enough people are unhappy about it && sees future in alternative (i.e. critical mass), we can collectively have a "phase transition".

Re: Supercharged high-resolution ocean simulation with Jax

#18

Post by civilized »

mjn wrote: post #14
adgjlsfhk1 wrote: post #9

Earlier quoted context omitted.

"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 b…

I read that as being about what language industry uses to write ML applications, not about technical feasibility of integrating machine learning methods into a codebase. Put differently: industry most often uses Python (especially in ML), therefore the author wants to target Python in order to maximize uptake outside of academia. They even admit that doing it in Python is technically harder than doing it in Julia ("U…

Somehow I don't think an ocean simulation needs to be in Python so some startup can use it to... what, sell ads or something?

Anyone interesting enough to be looking at your ocean simulation code can probably handle it being in Julia, and may even prefer it, since the language is so much better designed for this kind of thing than Python.

Re: Supercharged high-resolution ocean simulation with Jax

#19

Post by fault1 »

moelf wrote: post #17
mjn wrote: post #14

Earlier quoted context omitted.

I read that as being about what language industry uses to write ML applications, not about technical feasibility of integrating machine learning methods into a codebase. Put differently: industry most often uses Python (especially in ML), therefore the author wants to target Python in order to maximize uptake outside of academia. They even admit that doing it in Python is technically harder than doing it in Julia ("U…

hopefully if enough people are unhappy about it && sees future in alternative (i.e. critical mass), we can collectively have a "phase transition".

It may take a while, however. 15-20 years ago, you kind of had to use Python on the sly in the scientific setting vs the incumbents (MATLAB, C++, Fortran). Julia seems to be in a similar phase.

That being said, Python does have some structural advantages since it positions itself as a universal glue. It's much easier to gain a critical mass in that regard vs a niche area like scientific or numerical computing. That being said, Julia is probably underrated in general purpose usage.

Re: Supercharged high-resolution ocean simulation with Jax

#20

Post by adgjlsfhk1 »

fault1 wrote: post #19
moelf wrote: post #17

Earlier quoted context omitted.

hopefully if enough people are unhappy about it && sees future in alternative (i.e. critical mass), we can collectively have a "phase transition".

It may take a while, however. 15-20 years ago, you kind of had to use Python on the sly in the scientific setting vs the incumbents (MATLAB, C++, Fortran). Julia seems to be in a similar phase. That being said, Python does have some structural advantages since it positions itself as a universal glue. It's much easier to gain a critical mass in that regard vs a niche area like scientific or numerical computing. That b…

I think Julia has a much better path to wide adoption for numerical computing/HPC. It is a much better language for package developers (you pretty much never have to go to a lower level language and everything can compose together with much less work). If you look at Julia and Python packages with similar functionality, the Julia one will typically be much more general and 1/10th the lines of code. This is a pretty powerful incentive for on-boarding package devs.
Post reply on HN