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

JAX – NumPy on the CPU, GPU, and TPU

jax.readthedocs.io

Post reply on HN
Print view View original

101–110 of 147 posts

Re: JAX – NumPy on the CPU, GPU, and TPU

#101

Post by jarym »

patrickkidger wrote: post #98
jarym wrote: post #88

Earlier quoted context omitted.

Yea exactly - Python for training, Java/.NET for inference at production. I looked at approaches like GRPC and things but my case is a bit more time-sensitive and the latency added by going over a network layer was too much. For now I'm happy with Pytorch->ONNX and then running the ONNX model directly. But as I said, that means I can't easily train using JAX :-(

You can do JAX->TF->ONNX I believe: https://github.com/patrick-kidger/equinox/blob/main/equinox/...

Ohh, I'll check that out!

Re: JAX – NumPy on the CPU, GPU, and TPU

#102

Post by patrickkidger »

brap wrote: post #100
patrickkidger wrote: post #97

Earlier quoted context omitted.

Indeed, it is much better to use jit(grad(f)) in general. Supporting the opposite composition is still useful in some edge cases -- for example when debugging, you want to step through a computation without jit, and simply not crash when differentiating any inner functions also decorated with jit.

Wouldn't you use jax.disable_jit() for that?

You can do that too! That will disable every JIT though. In practice you might only want to disable just one.

To add some colour to my answer. When writing a library, it's typical to a put a JIT statement on everything in the public API. This means you get the benefits of JIT compilation even when you're just hacking around in the REPL, and mitigates the new-user-footgun in which they forget to use JIT themselves.

Meanwhile, good practice is always to JIT your whole computation.

Combined, this mean that it's fairly common to go jit (at the top level) -> grad (of your operation) -> jit (of some library call).

When debugging your code, the JIT'd library call is _probably_ not the culprit. So you only want to disable the top-level JIT when stepping through, and still take advantage of JIT compilation where you can. Overall one obtains a composition of the form grad(jit(...)).

TL;DR: even if use case doesn't come up super frequently, it's more user-friendly to support grad(jit(...)) than it is to just crash.

Re: JAX – NumPy on the CPU, GPU, and TPU

#103

Post by catgary »

cl3misch wrote: post #89
marmaduke wrote: post #28

Earlier quoted context omitted.

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.

Re: JAX – NumPy on the CPU, GPU, and TPU

#104

Post by cl3misch »

patrickkidger wrote: post #95
cl3misch wrote: post #46

Earlier quoted context omitted.

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

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.

Re: JAX – NumPy on the CPU, GPU, and TPU

#105

Post by LaszloKv »

dan_mctree wrote: post #72
antognini wrote: post #8

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

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)

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

Re: JAX – NumPy on the CPU, GPU, and TPU

#106

Post by cl3misch »

catgary wrote: post #103
cl3misch wrote: post #89

Earlier quoted context omitted.

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?

Re: JAX – NumPy on the CPU, GPU, and TPU

#107

Post by aportnoy »

amir734jj wrote: post #23
aportnoy wrote: post #19

Earlier quoted context omitted.

What happened?

Nothing happened. It was an informal technical interview with the program manager at JAX. 1 hour call and the interview was remote but describing him as opinionated and entitled is an understatement. Best of luck to them.

Asking because I am on that team :)

I went through the same hiring process and had a positive experience at every stage. I had a strong competing offer but went with the JAX team at NVIDIA.

I'll pass it along as feedback.

Re: JAX – NumPy on the CPU, GPU, and TPU

#108

Post by Q6T46nT668w6i3m »

antognini wrote: post #8

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

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.

Re: JAX – NumPy on the CPU, GPU, and TPU

#109

Post by Q6T46nT668w6i3m »

dan_mctree wrote: post #72
antognini wrote: post #8

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

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)

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

Re: JAX – NumPy on the CPU, GPU, and TPU

#110

Post by emu »

marmaduke wrote: post #30
subhrm wrote: post #7

JAX GPU support is limited to Linux only. Even the WSL2 support is experimental. https://jax.readthedocs.io/en/latest/installation.html#suppo...

They don't build on Windows at all, as well.

Not true!

We release Windows CPU wheels (https://pypi.org/project/jaxlib/#files). So JAX on CPU works great on Windows.

We don't release Windows GPU wheels at the moment, but that's because we're a small team and none of us use Windows personally. We welcome contributions!

(I verified that the Windows CUDA GPU support built as recently as two weeks ago, but I don't have the ability to test that it works.)

We recommend WSL2 because that's just using our existing Linux CUDA release.

Post reply on HN