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

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

#131

Post by Evidlo » Fri, Sep 29, 2023, 6:05 PM UTC

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…

Why is that? Why doesn't Jax just do something like

    class JaxWrapper:
        def __init__(self, arr):
            self.arr = arr
        def __setitem__(self, key, val):
            return self.arr.at[key].set(val)
        ....

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

#132

Post by tehsauce » Fri, Sep 29, 2023, 6:10 PM UTC

mattjjatgoogle wrote: post #112

tehsauce wrote: post #12

Earlier quoted context omitted.

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

I've seen it although haven't dived deep into it. It looks like they have some interesting support for particle cell data structures, but is fairly complicated and carries limitations: https://jax-md.readthedocs.io/en/main/_modules/jax_md/partit...

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

#133

Post by gorenb » Fri, Sep 29, 2023, 6:14 PM UTC

This is pretty well known, so I don't know why it would get so popular.

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

#134

Post by robertnishihara » Fri, Sep 29, 2023, 6:19 PM UTC

I'm a huge fan of Jax. The Jax team is incredibly strong!

Just want to share that Ray (an open source project we're developing at Anyscale), can be used to scale Jax (e.g., across TPUs).

Some docs from Google on how to do this

https://cloud.google.com/tpu/docs/ray-guide

Alpa is an open source project scaling Jax on 1000+ GPUs

https://www.anyscale.com/blog/training-175b-parameter-langua...

Cohere uses Ray + Jax + TPUs to build their LLMs

https://www.youtube.com/watch?v=For8yLkZP5w

A demo from Matt Johnson on the Jax team

https://www.youtube.com/watch?v=hyQ-tgD5sgc

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

#135

Post by dekhn » Fri, Sep 29, 2023, 6:39 PM UTC

I'll remember this always: when DeepMind solved a subset of the protein structure prediction problem, they used Jax as the framework.

PSPP was a long-standing issue and to see a fairly new computational tool used to significantly aid in the process of "solving" it speaks greatly towards its general utility in the sciences.

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

#136

Post by jjtheblunt » Fri, Sep 29, 2023, 6:42 PM UTC

mi_lk wrote: post #84

Looking for some numerical problems to practice JAX, any suggestion or resource?

there's a comment above : search for kaggle and you'll find such.

also:

https://jax.readthedocs.io/en/latest/advanced_guide.html

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

#137

Post by dekhn » Fri, Sep 29, 2023, 6:42 PM UTC

mattjjatgoogle wrote: post #112

tehsauce wrote: post #12

Earlier quoted context omitted.

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

Last time I looked at JAX MD it didn't support most of the force field terms necessary for simulating proteins and DNA. For example, it could do n-body simulations with some potential, but not the bonds/torsions between atoms. It's unclear if they added support, but that's a huge gap in functionality compared to other systems.

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

#138

Post by patrickkidger » Fri, Sep 29, 2023, 7:30 PM UTC

cl3misch wrote: post #119

patrickkidger wrote: post #113

Earlier quoted context omitted.

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

Ah right! Actually it's a good point, the Equinox readme/etc do tend to emphasise the ML use cases -- partly this is deliberate (go where the money is)! But I should probably tweak it to emphasise more general parameterised models.

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

#139

Post by marmaduke » Fri, Sep 29, 2023, 7:53 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?

there's a paper about it that I just found, enjoy https://arxiv.org/pdf/2301.13062

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

#140

Post by marmaduke » Fri, Sep 29, 2023, 7:58 PM UTC

emu wrote: post #129

marmaduke wrote: post #120

Earlier quoted context omitted.

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.

Very nice! I just installed, I hope to eventually contribute down the line, especially in terms of custom operators. They weren't even document until recently, and there's still quite some work to add them.