Live data from Hacker News

JAX – NumPy on the CPU, GPU, and TPU

jax.readthedocs.io

51–60 of 147 posts

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

#51
Imho the project should emphasize more the potential for simple and uniquitous multi-core acceleration of vector compute that is available by definition to anybody having any modern cpu.

nvidia dropped cuda support for perfectly good gpu's, showing the perils and waste of being locked-in in a profit-maximazing monopoly.

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

#52

Earlier quoted context omitted.

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.

It's my understanding that, at least in Python, you can't change immutable data type but you can just assign a new data to the same variable and therefore overwrite it, right? So even if JAX makes list type immutable, you can still just re-use `x` to save the new modified list.

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

#53

Is there any benefit using it instead of pytorch?

Jax has a much nicer handling of higher order differentiation. PyTorch has functions to compute Hessians and there are libraries to keep differentiability through optimizers, but going out of their standard use-cases becomes tricky very fast. In contrast, JAX can compute nth-derivatives of things very easily.

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

#56

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'll run on a GPU for free and be about 1000x faster.

Are there any benchmarks for that? Running on GPU never comes for free. You have to transfer data back and forth which has a cost, for instance.

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

#57

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

Meanwhile, the first sentence in their readme is this:

> JAX is Autograd and XLA, brought together for high-performance machine learning research.

That does not really convey the generality of it that well.

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

#58
post #20

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…

I found this to be not true in practice when working with graphs.

Having access to high performance explicit loops and ifs/masks allows one to focus on the hard parts of the algorithms, rather than on the purely incidental puzzle how to best avoid spending time in the Python runtime.

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

#60
post #36

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

The open source release of XLA predates Lattner's tenure at Google by 7 months, and it definitely existed before that -- the codebase was already 66k SLOC at that point. During his tenure it went from 100k SLOC to 250k SLOC. It's now 700k SLOC. He also has, as far as I can tell, zero commits in the XLA codebase. "Of LLVM fame" would be more accurate I think.
Post reply on HN