JAX – NumPy on the CPU, GPU, and TPU • Hacker News Board

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

#11

Post by adgjlsfhk1 » Fri, Sep 29, 2023, 4:44 AM UTC

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…

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)

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

#12

Post by tehsauce » Fri, Sep 29, 2023, 4:44 AM UTC

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…

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.

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

#13

Post by krasin » Fri, Sep 29, 2023, 4:54 AM UTC

adgjlsfhk1 wrote: post #11

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…

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.

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

#14

Post by nharada » Fri, Sep 29, 2023, 5:07 AM UTC

For large scale ML Jax is pretty nice. It makes the multi-host computation feel like a first class citizen and your programs tend to be designed with that in mind.

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

#15

Post by amir734jj » Fri, Sep 29, 2023, 5:15 AM UTC

Very unrelated but I did job interview with Nvidia JAX team for a compiler engineer role some time ago, not very friendly and very opinionated.

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

#16

Post by tasubotadas » Fri, Sep 29, 2023, 5:16 AM UTC

Is there any benefit using it instead of pytorch?

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

#17

Post by fnbr » Fri, Sep 29, 2023, 5:18 AM UTC

tasubotadas wrote: post #16

Is there any benefit using it instead of pytorch?

The main benefit in my experience is that it’s much easier to do distributed computations in JAX. It has a much nicer API. For single device computing there’s no advantage either way.

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

#18

Post by fnbr » Fri, Sep 29, 2023, 5:19 AM UTC

esafak wrote: post #4

Anybody using it in production? Is it, or its derivatives like Flax, worth using over pyTorch for anything? edit: Made comparison more fair.

A number of large AI companies use it to train their large models; Midjourney, Stability, Anthropic, DeepMind, among others.

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

#19

Post by aportnoy » Fri, Sep 29, 2023, 5:23 AM UTC

amir734jj wrote: post #15

Very unrelated but I did job interview with Nvidia JAX team for a compiler engineer role some time ago, not very friendly and very opinionated.

What happened?

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

#20

Post by nestorD » Fri, Sep 29, 2023, 5:29 AM UTC

adgjlsfhk1 wrote: post #11

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…

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)

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.