News for hackers. Stuff that matters.
Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by cl3misch on Friday September 29, 2023 @03:13AM (#37700294)
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.

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by patrickkidger on Friday September 29, 2023 @10:22AM (#37704711)
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.)

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by cl3misch on Friday September 29, 2023 @11:22AM (#37705579)
> 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.

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by patrickkidger on Friday September 29, 2023 @11:57AM (#37706114)
> 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.

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by cl3misch on Friday September 29, 2023 @01:01PM (#37707040)
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.

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by patrickkidger on Friday September 29, 2023 @03:30PM (#37709056)
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.
Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by ssivark on Friday September 29, 2023 @12:31PM (#37706613)
> 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!

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by cl3misch on Friday September 29, 2023 @01:08PM (#37707171)
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