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

31–40 of 147 posts

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

#31

Post by sdenton4 »

nestorD wrote: post #20
adgjlsfhk1 wrote: post #11

Earlier quoted context omitted.

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]: `ar…

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.

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

#33

Post by shihab »

nestorD wrote: post #20
adgjlsfhk1 wrote: post #11

Earlier quoted context omitted.

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]: `ar…

> 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

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

#34

Post by technocratius »

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.

We recently switched to Jax to boost performance as we scale up our algorithmic core. The nice thing is that it presents only a minor jump in capabilities to get developers working with it, if they have prior exposure to numpy ofcourse. Quite nice :)

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

#36

Post by swyx »

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…

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

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

#37

Post by jphoward »

I think JAX is cool, but I do find it slightly disingenuous when it claims to be "numpy by on the GPU" (as opposed to PyTorch), when actually there's a fundamental difference; it's functional. So if I have an array `x` and want to set index 0 to 10, I can't do:

  x[0] = 10
Instead I have to do:

  y = x.at[0].set(10)
Of course this has advantages, but you can't then go and claim that JAX is a drop in replacement for numpy, because this such a fundamental change to how numpy developers think (and in this regard, PyTorch is closer to numpy than JAX).

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

#38

Post by fulafel »

shihab wrote: post #33
nestorD wrote: post #20

Earlier quoted context omitted.

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]: `ar…

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

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.

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

#39

Post by LightFog »

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…

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.

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

#40

Post by j7ake »

jphoward wrote: post #37

I think JAX is cool, but I do find it slightly disingenuous when it claims to be "numpy by on the GPU" (as opposed to PyTorch), when actually there's a fundamental difference; it's functional. So if I have an array `x` and want to set index 0 to 10, I can't do: x[0] = 10 Instead I have to do: y = x.at[0].set(10) Of course this has advantages, but you can't then go and claim that JAX is a drop in replacement for numpy…

Also conditionals can be tricky (greater, if else) and often need rewriting.
Post reply on HN