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

141–147 of 147 posts

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

#141

Post by VHRanger »

woadwarrior01 wrote: post #91
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...

Apple supports JAX[0] along with PyTorch[1] and Tensorflow[2] on macOS with both Apple Silicon and AMD GPUs (on x86 Macs). Although, the perf isn't great. I write most of my experimental ML code in JAX on an M2 Macbook Air and then move to a proper multi-GPU Linux box for full training runs. [0]: https://developer.apple.com/metal/jax/ [1]: https://developer.apple.com/metal/pytorch/ [2]: https://developer.apple.com/me…

Pytorch on my M2 max using the MPS backend has pretty decent performance to be honest?

It's significantly faster than CPU. Something like 100x using sheet

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

#142

Post by wodenokoto »

maurosilber wrote: post #66
wodenokoto wrote: post #64

Earlier quoted context omitted.

doesn't `[] =` just call a method on the object in python? e.g, `x[0] = 10` is the same as `x.__set_item__(0, 10)`, so there shouldn't be any technical limitation to using `x[0]` (says the guy who never even imported jax)

You could do `y = x.__setitem__(0, 10)`, but you cannot assign `x[0] = 10` to a new variable. If `__setitem__` was overridden, you would not be able to distinguish between these cases and raise an error in the second one.

Yes, that makes perfect sense.

I somehow completely missed the assignment part of the second example.

Thank you for the clarification.

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

#143

Post by sheepscreek »

cl3misch wrote: post #85
sheepscreek wrote: post #82

I’ve used Numpy with Numba primarily (on the CPU) and it had been a game changer for my data science workloads. Naturally, I pay close attention to Jax and give it a look every now and then. So, I’ll focus my observations below on Jax’s Numpy API support. At a glance, Jax code looks like regular Python, but it’s a very different style of programming. Two big differences I’ve found are: - All Jax functions must be pur…

> You have to hardcode the shape tuples. This is not true. Rather, all shapes have to be known at compile time. That means that output shapes must not depend on input values , but may depend on input shapes -- also explicitely. Furthermore, there are two useful additions: 1. You can use vanilla numpy for compile-time computations. An example would be computing an array of indices for some moving-window filter, depend…

You’re right, I apologize for wording it incorrectly. It might have been a restriction with the # of dimensions. Either way, it wasn’t cut out for my use case.

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

#144

Post by pjmlp »

KeplerBoy wrote: post #50
pjtr wrote: post #22

Earlier quoted context omitted.

Is there a specific reason why Windows is not supported?

Presumably because the Google cloud doesn't run on Windows. Well, nothin HPC related runs Windows.

Life science industry uses plenty of Windows, including HPC workloads.

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

#145

Post by mattjjatgoogle »

cl3misch wrote: post #87
sampo wrote: post #68

> With its updated version of Autograd, JAX can automatically differentiate native Python and NumPy code. Nice. When did they make this change? Here is the old way in the docs, where you needed to define functions for the if-true branch and the if-false branch, and feed them to a conditional function, to get the normal if-then-else conditional. https://jax.readthedocs.io/en/latest/notebooks/Common_Gotcha...

This is still the case afaik. For vanilla "if", the condition must be known at compile time. For runtime, you have to use "cond", "where", or "select" (which may be analogous).

Actually, that's never been a constraint for JAX autodiff. JAX grew out of the original Autograd (https://github.com/hips/autograd), so differentiating through Python control flow always worked. It's jax.jit and jax.vmap which place constraints on control flow, requiring structured control flow combinators like those.

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

#146

Post by mattjjatgoogle »

sampo wrote: post #68

> With its updated version of Autograd, JAX can automatically differentiate native Python and NumPy code. Nice. When did they make this change? Here is the old way in the docs, where you needed to define functions for the if-true branch and the if-false branch, and feed them to a conditional function, to get the normal if-then-else conditional. https://jax.readthedocs.io/en/latest/notebooks/Common_Gotcha...

Actually that never changed. The README has always had an example of differentiating through native Python control flow:

https://github.com/google/jax/commit/948a8db0adf233f333f3e5f...

The constraints on control flow expressions come from jax.jit (because Python control flow can't be staged out) and jax.vmap (because we can't take multiple branches of Python control flow, which we might need to do for different batch elements). But autodiff of Python-native control flow works fine!

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

#147

Post by [deleted] »

thrdbndndn wrote: post #41
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…

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…

[deleted]
Post reply on HN