Hacker News Reader

Comment by antognini | Hacker News Reader

It took me a while to realize it, but Jax is actually a huge opportunity for a lot of scientific computing. Jax was originally developed as a more flexible platform for doing machine learning research. But Jax's real superpower is that it bundles XLA and makes it really easy to run computations on GPU or TPU. And huge swathes of scientific computation basically run large scale vectorized computations.

When I was in astronomy (about a decade ago) I did large scale simulations of gravitational interactions. But at the time all these simulations were done on CPU. Some of the really big efforts used more specialized chips, but it was a huge effort to write the code for it.

But today with Jax, if you want to write an N-body simulation of a globular cluster, you can just code it up in numpy and it'll run on a GPU for free and be about 1000x faster. From what I can tell though, very few people in the sciences have caught on yet.

View on HN

Luckily I discovered JAX in the beginning of my PhD four years ago. It has made our data processing (biomedical imaging) so much easier and more readable, albeit with a slight learning curve due to JAX being functional/pure.

I am also continously surprised how little adoption JIT and autodiff libraries have gotten in scientific computing. A lot of my colleagues somehow really like coding cost function gradients and fine-tuned GPU code by hand. I guess using something like JAX can reduce your standing in the group, because it can make it seem like coding algorithms is pretty easy.

I feel like I see the opposite -- that everything scientific computing is getting rewritten in something autodifferentiable! Whether that's JAX or something else.

My experience might be biased though: shameless advert for Equinox (https://github.com/patrick-kidger/equinox, 1.4k GitHub stars), which is now the foundation of quite a lot of SciComp in JAX. (Both internal and open-source.)

> foundation of quite a lot of SciComp in JAX

...if that SciComp uses machine learning, I guess? In my "physics of biomedical imaging" bubble, people are hardly doing state-of-the-art ML, but rather expensive forward models for which computing a gradient is cumbersome.

But I know that e.g. Stephan Hoyer is a physicist and you are a mathematician originally -- I have read a lot of your JAX issues and libraries ;-) maybe it just depends on the "mini-bubble' aka. the indiviual research group and not only the field of science.

> if that SciComp uses machine learning, I guess?

Not necessarily! It's perfectly possible (and quite common) to e.g. write down a traditional parameterised ODE, and then optimise its parameters via gradient descent. Compute the gradients wrt parameters using autodiff through the numerical ODE solver. All without a single neural network in sight! ;)

My usual spiel is that autodiff+autoparallel are really useful for any kind of numerical computation -- of which ML is a (popular, well funded) special case.

At least in my mini bubble, these kinds of "scipy but autodifferentiable" use-cases are fairly common.

> I have read a lot of your JAX issues and libraries ;-)

Haha, that's fun to hear though! Thank you for sharing that.

Ah, I think I was unclear. I specifically meant your reference to Equinox, because that seemed to me to be somewhat ML specific.

In general, I very much agree that "autodiff+autoparallel are really useful for any kind of numerical computation". And the use cases are also really common in my bubble. It's just that (imho) most people have not realized this.

Ah right! Actually it's a good point, the Equinox readme/etc do tend to emphasise the ML use cases -- partly this is deliberate (go where the money is)! But I should probably tweak it to emphasise more general parameterised models.

> In my "physics of biomedical imaging" bubble, people are hardly doing state-of-the-art ML, but rather expensive forward models for which computing a gradient is cumbersome.

I’d appreciate any pointers to the literature; curious to see the kinds of models people work with. Thanks!

I don't have didactic examples at hand, but e.g. [1] or [2]. IIRC [1] uses the Laplace operator (second-order spatial derivate) and [2] uses a linear solve inside the forward model through which differentiation is certainly possible but pretty cumbersome in practice.

[1] https://www.nature.com/articles/s41598-019-52283-6

[2] https://doi.org/10.1117/1.JMI.4.3.034005

> "It took me a while to realize it, but Jax is actually a huge opportunity for a lot of scientific computing."

In all conferences like NeurIPS, in Google ML Community days, etc., whenever there is a JAX workshop/tutorial/talk, it is always touted as a numerical computation library. And it was developed as such. Sure the focus is in ML, but everyone involved in it always have said that this is a general purpose scientific computing library.

Flax, Haiku, etc. are Deep Learning libraries.

Meanwhile, the first sentence in their readme is this:

> JAX is Autograd and XLA, brought together for high-performance machine learning research.

That does not really convey the generality of it that well.

You're right! Maybe we should revise that... I made https://github.com/google/jax/pull/17851, comments welcome!

So was tensorflow... And yet it's pretty much dead.

None of these libraries unfortunately allows making good use of CPU vectorized units. Xla might produce some SIMD code but it pales in (performance) comparison to routines written explicitly for SIMD on GPU. ISPC is a good example of this.

the issue here is that if your ideal algorithm isn't simply expressible in numpy (which many aren't), you're pretty much out of luck. As a result, imo the better approach is to use a fast language that also compiles to GPU (e.g. Julia)

An alternative is to write most of the program in Python + JAX + implement a few custom XLA ops in CUDA / Triton. That way, the program is very readable and can interoperate with the larger ecosystem, while still being fast to run.

Having used JAX quite a bit for numerical computing (and having lectured on this use-case) I would say that a surprisingly large number of algorithms can be expressed as array[0] operations (even if it sometimes takes a bit of thinking).

And, more importantly, things that cannot be expressed that way tend to not be a good fit for GPU computing anyway (independently of the language / framework you are using).

[0]: `array` is a shortcut here, JAX is not limited to operations on arrays.

Agreed. I've done a fair amount of reworking signal processing algorithms to run on GPU/TPU, and it's a different beast. You often have to really rebuild the algorithm from the ground up to take advantage of parallelization. But often you /can/ rework the algorithm, and end up with much higher throughput than the crusty old serial algorithm: there's typically nothing fundamentally stopping you from finding a good implementation, just that the original devs were working in the 70s and hasn't thought that far ahead.

> things that cannot be expressed that way tend to not be a good fit for GPU computing anyway

I'll have to disagree with you a little bit here. SIMT model of GPUs are quiet a bit more expressive than the numpy's SIMD model. As an obvious example, you'll have to manually maintain a mask to implement if/else i.e. code path divergence in SIMD. GPUs automatically does this and many more to make your life easier. And frankly, I find it lot more easier to reason about what should happen to one data point than a bunch of them together.

An interesting article I read recently that has some relevance to this discussion. https://pharr.org/matt/blog/2018/04/18/ispc-origins

To spell out what the linked ISPC post implies, most of the difference, like ISPC shows, is differences in GPU languages and compilers vs CPU side equivalents.

Convenient authoring doesn't necessarily make it a good fit for the hardware. Add in enough divergence and your GPU code is going to be matched or outperformed by a competent CPU implementation (on a chip of comparable size). Branchless code can result in substantial speedups on either.

To be fair though, modern GPUs are pretty good at branching and latency hiding, while numpy-style code has poor data locality unless you have a magic compiler.

I found this to be not true in practice when working with graphs.

Having access to high performance explicit loops and ifs/masks allows one to focus on the hard parts of the algorithms, rather than on the purely incidental puzzle how to best avoid spending time in the Python runtime.

Pretty sure I got multiple 1000x speed ups when I vectorized my my algo trader from a dumb python loop to a dumb numba compiled thing, and when I benchmarked Jax, the performance blew away the numba thing (which was already a million times faster than the naive version) because Jax performance stayed perfectly flat as the scale went up whereas numba slowed down. Might have been my approach for each, but it was enlightening and funny to watch.

I have at least one complaint with the numpy model:

When you chain a sequence of vectorized operations on arrays, loop fusion would save you from allocating memory for each intermediate variable, and the round trip time of moving it from RAM to CPU multiple times. I don’t know how good JAX’s JITted loop fusion is on CPU, but I’ve been very very impressed by Julia.

Eg: I had some Numpy code that took hours (and needed terabyte RAM) that was very straightforward to code in Julia, and needed only a few GB to finish in a few seconds — on my laptop.

I want to be able to think in arrays, but to also not have to materialize the arrays as much as possible.

Jax JIT of scan is fairly good, so loops aren't as slow as you'd expect.

A key difference is that each iteration of scan is called by the host. Put differently, JAX can't fuse scan into a single GPU kernel, but launches a kernel for each iteration.

Depening on the workload this is no problem. If you have many cheap iterations, you will notice the overhead.

I am not sure if they are working on fusing scan and what's the current status.

There is an ‘unroll’ parameter in scan that lets you control how many iterations of the loop are fused into a single kernel.

Yes, but is it really the same? Afaik the `unroll=n` parameter translates `n` iterations into a vanilla `for` loop which is then unrolled into sequential statements (in contrast to a JAX `fori` loop). There still is no loop on the accelerator, strictly speaking?

I think this is up to XLA to handle not Jax. The whole selling point in TF of the tf.function decorator (which uses XLA underneath as well) is that it fuses arithmetic to lower launch count.

there's a paper about it that I just found, enjoy https://arxiv.org/pdf/2301.13062

Jax is super useful for scientific computing. Although nbody sims might not be the best application. A naive nbody sim is very easy to implement and accelerate in jax (here’s my version: https://github.com/PWhiddy/jax-experiments/blob/main/nbody.i...), but it can be tricky to scale it. This is because efficient nbody sims usually either rely on trees or spatial hashing/sorting which are tricky to efficiently implement with jax.

Have you seen JAX MD? https://github.com/jax-md/jax-md

I've seen it although haven't dived deep into it. It looks like they have some interesting support for particle cell data structures, but is fairly complicated and carries limitations: https://jax-md.readthedocs.io/en/main/_modules/jax_md/partit...

Last time I looked at JAX MD it didn't support most of the force field terms necessary for simulating proteins and DNA. For example, it could do n-body simulations with some potential, but not the bonds/torsions between atoms. It's unclear if they added support, but that's a huge gap in functionality compared to other systems.

we did an interview with Chris Lattner of XLA fame where he also similarly had nice things to say about JAX: https://www.latent.space/p/modular

just sharing for those who want to learn more

The open source release of XLA predates Lattner's tenure at Google by 7 months, and it definitely existed before that -- the codebase was already 66k SLOC at that point. During his tenure it went from 100k SLOC to 250k SLOC. It's now 700k SLOC. He also has, as far as I can tell, zero commits in the XLA codebase. "Of LLVM fame" would be more accurate I think.

my bad - i guess i was just saying he led that team but didnt mean to imply he originated it

you seem to have very precise knowledge of the SLOC at a point in time - just curious is there any tooling you used to do that? that can be pretty nifty to pull out on occasion

I git cloned the repo and then ran sloccount after checking out various commits (just did `git log | grep -C3 'Jan 1 [0-9:]* 2017'` or similar to find the relevant commits)

ha, simple enough. thx

> large scale simulations of gravitational interactions

I'm guessing this was mostly Fast Multipole Method? I don't think it ports that easily GPU since there's so much communication involved and the leaves don't do a whole lot

Yes. But some of the algorithms cannot benefit that much from the GPU. In my field -- mathematical optimization, lots of algorithms rely on sparse matrix operations and takes many iterations until convergence.

Would this help? https://jax.readthedocs.io/en/latest/jax.experimental.sparse...

Nope it's super slow for large sparse matrices. It's even faster to use generic scatter/gather to implement some, instead of that built in thing.

Have you investigated why? I know that many projects have an "implement first, optimize later" approach, and the lesser used functions might be far from optimal.

Back in the tensorflow days, I had this issue and submitted a patch that gave a ~50x speedup for my usecase. It's always better to optimize the base function rather than have 100 people all manually working around the same performance issue.

Because they use a funny format (BCOO). I'm not mocking, it must be a solid choice for some reasons, like sparsification or other fancy stuff. But for large and even with batches (ie multiply with tall dense matrix), it doesn't match an equivalent scatter (x.at[idx].add(vals)). Which itself is several times slow than equivalent opencl (on an A40)

> it'll run on a GPU for free and be about 1000x faster.

Are there any benchmarks for that? Running on GPU never comes for free. You have to transfer data back and forth which has a cost, for instance.

That was just from some quick benchmarks I did a few months back on some 10,000 particle N-body simulations. The performance boost will depend on the task, though. For the kinds of computations I did in grad school it would have been less effective since I was only looking at 3--5 objects, so there's just less parallelism to take advantage of.

Would you happen to have sources on the three orders of magnitude speedup coming at no cost? I'd assume porting + data movement considerations making this task non-trivial.

Yep. I’ve been writing a molecular dynamics simulator in PyTorch and the reduction in scope is wild because of heterogenous operators and automatic differentiation.

Porting existing codes is still a massive effort and there is low faith in long-term support from Google based software and hardware. I’m not aware on much (any) TPU use in scientific HPC.

There is not yet. But there is huge pressure on HPC Centers at least in Europe to also make resources available for ML. Already many scientific supercomputers have GPUs as accelerators. So we might have the other situation: HPC users are faced more and more with machines with spare accelerators and it will make sense to use them. Actually it would totally make sense if JAX development is in part also public also financed in this case (e.g. through EuroHPC).

It's so awkward that these truly fantastic tools for fast numerical computation (NumPy and JAX) have to be accessed through Python, which is a truly terrible language for fast numerical computation.

Is anyone making any serious progress in fast GPU based computational tools for other faster languages? I'm looking for something that also works on the GPU on windows (unlike JAX)

> have to be accessed through python

It’s because most of the people doing these computations don’t have the capacity to become experts in multiple fields. They understand the math and analytics very well, and they expend all their time thinking about that, not about type systems, memory management, etc. Python lets them code without having to think about a lot of that stuff so they can focus on the things they care about. These aren’t computer scientists or programmers, they’re meteorologists, astronomers, oil and gas analysts, investment bankers etc. That’s why some truly great computer scientists and programmers invested their time into building these tools for python vs other languages.

You might want to look into Futhark https://futhark-lang.org/

This type of code isn’t executed by the Python interpreter.

Reply on news.ycombinator.com