JAX – NumPy on the CPU, GPU, and TPU • Hacker News Board

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

#121

Post by cl3misch » Fri, Sep 29, 2023, 5:08 PM UTC

ssivark wrote: post #115

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…

> 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!

I don't have didactic examples at hand, but e.g. [1] or [2]. IIRC [1] uses the Laplace operator (second-order spatial derivate) and [2] uses a linear solve inside the forward model through which differentiation is certainly possible but pretty cumbersome in practice.

[1] https://www.nature.com/articles/s41598-019-52283-6

[2] https://doi.org/10.1117/1.JMI.4.3.034005

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

#122

Post by marmaduke » Fri, Sep 29, 2023, 5:11 PM UTC

cl3misch wrote: post #106

catgary wrote: post #103

Earlier quoted context omitted.

There is an ‘unroll’ parameter in scan that lets you control how many iterations of the loop are fused into a single kernel.

Yes, but is it really the same? Afaik the `unroll=n` parameter translates `n` iterations into a vanilla `for` loop which is then unrolled into sequential statements (in contrast to a JAX `fori` loop). There still is no loop on the accelerator, strictly speaking?

I think this is up to XLA to handle not Jax. The whole selling point in TF of the tf.function decorator (which uses XLA underneath as well) is that it fuses arithmetic to lower launch count.

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

#123

Post by marmaduke » Fri, Sep 29, 2023, 5:15 PM UTC

londons_explore wrote: post #74

marmaduke wrote: post #27

Earlier quoted context omitted.

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.

Because they use a funny format (BCOO). I'm not mocking, it must be a solid choice for some reasons, like sparsification or other fancy stuff. But for large and even with batches (ie multiply with tall dense matrix), it doesn't match an equivalent scatter (x.at[idx].add(vals)). Which itself is several times slow than equivalent opencl (on an A40)

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

#124

Post by joey_bob » Fri, Sep 29, 2023, 5:20 PM UTC

I honestly had no idea JAX was useful outside autograd, due to my own tunnel vision. I even used it with other libraries to do this kind of work. Is there a term for this type of mistake?

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

#125

Post by ayakang31415 » Fri, Sep 29, 2023, 5:33 PM UTC

polygamous_bat wrote: post #90

ayakang31415 wrote: post #21

I have been working on my DNN model using TensorFlow even though ML is not my main research. But it is a substantial part of my research, so I have to figure things out on my own, and I have done so over the past 3 years. However, I spend so much time on figuring out how any of TF methods works and debugging them. I never used JAX but I am not sure if this sort of grinding is normal when you use JAX as well (I always…

Is there a reason you are limited to JAX/tensorflow and can't use pytorch? For a lot of people I know whose main job was not to write code, switching from tensorflow to pytorch was something that saved them ten to hundreds of hours in the long run, even accounting for the initial learning time.

I always feel like I am a prisoner of sunk cost fallacy

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

#126

Post by rg111 » Fri, Sep 29, 2023, 5:43 PM UTC

ayakang31415 wrote: post #125

polygamous_bat wrote: post #90

Earlier quoted context omitted.

Is there a reason you are limited to JAX/tensorflow and can't use pytorch? For a lot of people I know whose main job was not to write code, switching from tensorflow to pytorch was something that saved them ten to hundreds of hours in the long run, even accounting for the initial learning time.

I always feel like I am a prisoner of sunk cost fallacy

You should definitely switch to PyTorch. Or even JAX.

TF is not worth it anymore.

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

#127

Post by __rito__ » Fri, Sep 29, 2023, 5:48 PM UTC

I wrote a Notebook [0] that gets you introduced to JAX in a very gentle manner.

It also covers things like functional purity in Deep Learning, and handling of random numbers in JAX.

[0]: Learn JAX: From Linear Regression to Neural Networks - https://www.kaggle.com/code/truthr/jax-0

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

#128

Post by __rito__ » Fri, Sep 29, 2023, 5:49 PM UTC

wjessup wrote: post #81

Great basics tutorial: https://colinraffel.com/blog/you-don-t-know-jax.html

Here's mine [0]. It's a very gentle introduction.

[0]: https://www.kaggle.com/code/truthr/jax-0

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

#129

Post by emu » Fri, Sep 29, 2023, 5:49 PM UTC

marmaduke wrote: post #120

emu wrote: post #110

Earlier quoted context omitted.

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?

Yes, we made this more formally supported recently.

We felt that Windows CPU support was important so everyone can run JAX, even if it's not always the most-accelerated version of JAX. And we got some great PRs from the community that helped fix a few open issues.

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

#130

Post by swyx » Fri, Sep 29, 2023, 5:52 PM UTC

taktoa wrote: post #117

swyx wrote: post #116

Earlier quoted context omitted.

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)

ha, simple enough. thx