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

41–50 of 147 posts

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

#41

Post by thrdbndndn »

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…

Sorry for my potentially VERY ignorant question, I only know functional programming at average joe level.

Why can't you do the first in functional programming (not in this specific case because it's just how it is, but in general)?

And even if you can't do so for any reasonable reason in functional (again, in general), what stops us to just add syntactic sugar to equal it to the second to make programmer's life easier?

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

#42

Post by microtonal »

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…

Agree, though I wouldn’t call PyTorch close to a drop-in for NumPy either, there are quite some mismatches in their APIs. CuPy is the drop-in. Excepting some corner cases, you can use the same code for both. E.g. Thinc’s ops work with both NumPy and CuPy:

https://github.com/explosion/thinc/blob/master/thinc/backend...

Though I guess the question is why one would still use NumPy when there are good libraries for CPU and GPU. Maybe for interop with other libraries, but DLPack works pretty well for converting arrays.

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

#43

Post by Hendrikto »

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.

Would you mind sharing some details? It sounds like an interesting peek behind the curtain.

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

#44

Post by riedel »

LightFog wrote: post #39
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.

There is not yet. But there is huge pressure on HPC Centers at least in Europe to also make resources available for ML. Already many scientific supercomputers have GPUs as accelerators. So we might have the other situation: HPC users are faced more and more with machines with spare accelerators and it will make sense to use them. Actually it would totally make sense if JAX development is in part also public also financed in this case (e.g. through EuroHPC).

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

#45

Post by Polygator »

thrdbndndn wrote: post #41
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…

Sorry for my potentially VERY ignorant question, I only know functional programming at average joe level. Why can't you do the first in functional programming (not in this specific case because it's just how it is, but in general)? And even if you can't do so for any reasonable reason in functional (again, in general), what stops us to just add syntactic sugar to equal it to the second to make programmer's life easie…

The fundamental reason why many functional languages won't allow you to do the first is that they use immutable data structures.

We could indeed introduce syntactic sugar (`y= (x[0]:=10)` maybe), but you'll still need to introduce a new variable to hold the modified list.

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

#46

Post by cl3misch »

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…

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

#47

Post by __rito__ »

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 took me a while to realize it, but Jax is actually a huge opportunity for a lot of scientific computing."

In all conferences like NeurIPS, in Google ML Community days, etc., whenever there is a JAX workshop/tutorial/talk, it is always touted as a numerical computation library. And it was developed as such. Sure the focus is in ML, but everyone involved in it always have said that this is a general purpose scientific computing library.

Flax, Haiku, etc. are Deep Learning libraries.

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

#48

Post by albertzeyer »

Does it support arrays of variable lengths now? Last time I looked, I think this was not supported. So it means, for every variable dimension, you need to use an upper bound, and then use masking properly, and hope that it would not waste computation too much on the unused part (e.g. when running a loop over it).

I'm working with sequences, e.g. speech recognition, machine translation, language modeling. This is a quite fundamental property for this type of models, that we have variable lengths sequences.

In those cases, for some example code, I have seen that training also used only fixed size dimensions. And at inference time, they had some non-JAX code for the loop over the sequence around the JAX code with fixed-size dimensions.

This seems like a quite fundamental issue to me? I wonder a bit that this is not an issue for others.

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

#50

Post by KeplerBoy »

pjtr wrote: post #22
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...

Is there a specific reason why Windows is not supported?

Presumably because the Google cloud doesn't run on Windows. Well, nothin HPC related runs Windows.
Post reply on HN