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

111–120 of 147 posts

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

#111

Post by emu »

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?

We ship Windows CPU only at the moment.

We don't support Windows GPU because we haven't had the engineer bandwidth to support it well.

We recommend WSL2 for GPU on Windows at the moment because that is a compromise: it allows CUDA support, without us having to support another release variant.

But we welcome community contributions!

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

#112

Post by mattjjatgoogle »

tehsauce wrote: post #12
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 impleme…

Have you seen JAX MD? https://github.com/jax-md/jax-md

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

#113

Post by patrickkidger »

cl3misch wrote: post #104
patrickkidger wrote: post #95

Earlier quoted context omitted.

I feel like I see the opposite -- that everything scientific computing is getting rewritten in something autodifferentiable! Whether that's JAX or something else. My experience might be biased though: shameless advert for Equinox ( https://github.com/patrick-kidger/equinox , 1.4k GitHub stars), which is now the foundation of quite a lot of SciComp in JAX. (Both internal and open-source.)

> foundation of quite a lot of SciComp in JAX ...if that SciComp uses machine learning, I guess? In my "physics of biomedical imaging" bubble, people are hardly doing state-of-the-art ML, but rather expensive forward models for which computing a gradient is cumbersome. But I know that e.g. Stephan Hoyer is a physicist and you are a mathematician originally -- I have read a lot of your JAX issues and libraries ;-) may…

> if that SciComp uses machine learning, I guess?

Not necessarily! It's perfectly possible (and quite common) to e.g. write down a traditional parameterised ODE, and then optimise its parameters via gradient descent. Compute the gradients wrt parameters using autodiff through the numerical ODE solver. All without a single neural network in sight! ;)

My usual spiel is that autodiff+autoparallel are really useful for any kind of numerical computation -- of which ML is a (popular, well funded) special case.

At least in my mini bubble, these kinds of "scipy but autodifferentiable" use-cases are fairly common.

> I have read a lot of your JAX issues and libraries ;-)

Haha, that's fun to hear though! Thank you for sharing that.

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

#114

Post by mattjjatgoogle »

matrss wrote: post #57
__rito__ wrote: post #47

Earlier quoted context omitted.

> "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.

You're right! Maybe we should revise that... I made https://github.com/google/jax/pull/17851, comments welcome!

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

#115

Post by ssivark »

cl3misch wrote: post #104
patrickkidger wrote: post #95

Earlier quoted context omitted.

I feel like I see the opposite -- that everything scientific computing is getting rewritten in something autodifferentiable! Whether that's JAX or something else. My experience might be biased though: shameless advert for Equinox ( https://github.com/patrick-kidger/equinox , 1.4k GitHub stars), which is now the foundation of quite a lot of SciComp in JAX. (Both internal and open-source.)

> foundation of quite a lot of SciComp in JAX ...if that SciComp uses machine learning, I guess? In my "physics of biomedical imaging" bubble, people are hardly doing state-of-the-art ML, but rather expensive forward models for which computing a gradient is cumbersome. But I know that e.g. Stephan Hoyer is a physicist and you are a mathematician originally -- I have read a lot of your JAX issues and libraries ;-) may…

> In my "physics of biomedical imaging" bubble, people are hardly doing state-of-the-art ML, but rather expensive forward models for which computing a gradient is cumbersome.

I’d appreciate any pointers to the literature; curious to see the kinds of models people work with. Thanks!

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

#116

Post by swyx »

taktoa wrote: post #60
swyx wrote: post #36

Earlier quoted context omitted.

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.

my bad - i guess i was just saying he led that team but didnt mean to imply he originated it

you seem to have very precise knowledge of the SLOC at a point in time - just curious is there any tooling you used to do that? that can be pretty nifty to pull out on occasion

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

#117

Post by taktoa »

swyx wrote: post #116
taktoa wrote: post #60

Earlier quoted context omitted.

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.

my bad - i guess i was just saying he led that team but didnt mean to imply he originated it you seem to have very precise knowledge of the SLOC at a point in time - just curious is there any tooling you used to do that? that can be pretty nifty to pull out on occasion

I git cloned the repo and then ran sloccount after checking out various commits (just did `git log | grep -C3 'Jan 1 [0-9:]* 2017'` or similar to find the relevant commits)

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

#118

Post by ssivark »

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…

I have at least one complaint with the numpy model:

When you chain a sequence of vectorized operations on arrays, loop fusion would save you from allocating memory for each intermediate variable, and the round trip time of moving it from RAM to CPU multiple times. I don’t know how good JAX’s JITted loop fusion is on CPU, but I’ve been very very impressed by Julia.

Eg: I had some Numpy code that took hours (and needed terabyte RAM) that was very straightforward to code in Julia, and needed only a few GB to finish in a few seconds — on my laptop.

I want to be able to think in arrays, but to also not have to materialize the arrays as much as possible.

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

#119

Post by cl3misch »

patrickkidger wrote: post #113
cl3misch wrote: post #104

Earlier quoted context omitted.

> foundation of quite a lot of SciComp in JAX ...if that SciComp uses machine learning, I guess? In my "physics of biomedical imaging" bubble, people are hardly doing state-of-the-art ML, but rather expensive forward models for which computing a gradient is cumbersome. But I know that e.g. Stephan Hoyer is a physicist and you are a mathematician originally -- I have read a lot of your JAX issues and libraries ;-) may…

> if that SciComp uses machine learning, I guess? Not necessarily! It's perfectly possible (and quite common) to e.g. write down a traditional parameterised ODE, and then optimise its parameters via gradient descent. Compute the gradients wrt parameters using autodiff through the numerical ODE solver. All without a single neural network in sight! ;) My usual spiel is that autodiff+autoparallel are really useful for a…

Ah, I think I was unclear. I specifically meant your reference to Equinox, because that seemed to me to be somewhat ML specific.

In general, I very much agree that "autodiff+autoparallel are really useful for any kind of numerical computation". And the use cases are also really common in my bubble. It's just that (imho) most people have not realized this.

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

#120

Post by marmaduke »

emu wrote: post #110
marmaduke wrote: post #30

Earlier quoted context omitted.

They don't build on Windows at all, as well.

Not true! We release Windows CPU wheels ( https://pypi.org/project/jaxlib/#files ). So JAX on CPU works great on Windows. We don't release Windows GPU wheels at the moment, but that's because we're a small team and none of us use Windows personally. We welcome contributions! (I verified that the Windows CUDA GPU support built as recently as two weeks ago, but I don't have the ability to test that it works.) We recomm…

Oops so sorry. But this is recent isn't it? I thought it was actually due to XLA/Bazel not supporting it?
Post reply on HN