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

71–80 of 147 posts

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

#71

Post by karterk »

albertzeyer wrote: post #48

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

For JIT-ing you need to know the sizes upfront. There was an experimental branch for introducing jagged tensors, but as far as I know, it has been abandoned.

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

#72

Post by dan_mctree »

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)

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

#73

Post by londons_explore »

__rito__ wrote: post #47
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 scien…

So was tensorflow... And yet it's pretty much dead.

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

#74

Post by londons_explore »

marmaduke wrote: post #27
modeless wrote: post #26

Earlier quoted context omitted.

Would this help? https://jax.readthedocs.io/en/latest/jax.experimental.sparse...

Nope it's super slow for large sparse matrices. It's even faster to use generic scatter/gather to implement some, instead of that built in thing.

Have you investigated why? I know that many projects have an "implement first, optimize later" approach, and the lesser used functions might be far from optimal.

Back in the tensorflow days, I had this issue and submitted a patch that gave a ~50x speedup for my usecase. It's always better to optimize the base function rather than have 100 people all manually working around the same performance issue.

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

#75

Post by rcpt »

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…

> large scale simulations of gravitational interactions

I'm guessing this was mostly Fast Multipole Method? I don't think it ports that easily GPU since there's so much communication involved and the leaves don't do a whole lot

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

#76

Post by brahbrah »

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)

> have to be accessed through python

It’s because most of the people doing these computations don’t have the capacity to become experts in multiple fields. They understand the math and analytics very well, and they expend all their time thinking about that, not about type systems, memory management, etc. Python lets them code without having to think about a lot of that stuff so they can focus on the things they care about. These aren’t computer scientists or programmers, they’re meteorologists, astronomers, oil and gas analysts, investment bankers etc. That’s why some truly great computer scientists and programmers invested their time into building these tools for python vs other languages.

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

#77

Post by bionhoward »

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…

Pretty sure I got multiple 1000x speed ups when I vectorized my my algo trader from a dumb python loop to a dumb numba compiled thing, and when I benchmarked Jax, the performance blew away the numba thing (which was already a million times faster than the naive version) because Jax performance stayed perfectly flat as the scale went up whereas numba slowed down. Might have been my approach for each, but it was enlightening and funny to watch.

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

#79

Post by varelse »

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…

[dead]
Post reply on HN