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.
> 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.
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).
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/
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?