Hacker News Reader

Comment by cl3misch | Hacker News Reader

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.

View on HN

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

Reply on news.ycombinator.com