Hacker News Board Hacker News, browsed like a classic message board
Live data from Hacker News

Jax and Equinox: What are they and why should I bother?

garymm.org

Post reply on HN
Print view View original

11–17 of 17 posts

Re: Jax and Equinox: What are they and why should I bother?

#11

Post by Archit3ch »

crazygringo wrote: post #6
Lerc wrote: post #4

For me the questions to answer for whether or not I should bother. Will it try and bind me to other technologies? Does it work out of the box on ${GPU}? Is it well supported? Will it continue to be supported?

Playing with JAX on Google Colab (Nvidia T4), everything works great. Sadly, I cannot get JAX to work with the built-in GPU on my M1 MacBook Air. In theory it's supposed to work: https://developer.apple.com/metal/jax/ But it crashes Python when I try to run a compiled function. And that's only after discovering I need an older specific version of jax-metal, because newer versions apparently don't work with M1 anymore…

Outside training, M1 also makes sense as a lower-bound deployment target ("we support Apple Silicon").

Re: Jax and Equinox: What are they and why should I bother?

#12

Post by __rito__ »

p1esk wrote: post #7

I wish Jax had everything I need to experiment with DL models built in, natively - like Pytorch. Instead there are many third party libraries (flax, trax, haiku, this one, etc). I have no idea which one to use. This was the case when I first played with jax 5 years ago, and it’s still the case today (even worse it seems). This makes it a non starter for me.

Why?

Use any. I used to work with Flax, now I work with Equinox more. Choose any between Flax, Equinox, and Haiku.

Re: Jax and Equinox: What are they and why should I bother?

#13

Post by spearman »

mccoyb wrote: post #10
spearman wrote: post #8

Earlier quoted context omitted.

Pretty sure it's Apple building it, and they're using JAX in-house so I imagine it will get better over time. Though they do love to drop support for old things so maybe M1 will never work again...

I think they’re likely using MLX in house now, no? (Probably not everyone, ofc - but seems likely that many will just use the native array framework designed explicitly for MX chips)

From https://machinelearning.apple.com/research/introducing-apple...

> Our foundation models are trained on Apple's AXLearn framework, an open-source project we released in 2023. It builds on top of JAX

Re: Jax and Equinox: What are they and why should I bother?

#14

Post by p1esk »

__rito__ wrote: post #12
p1esk wrote: post #7

I wish Jax had everything I need to experiment with DL models built in, natively - like Pytorch. Instead there are many third party libraries (flax, trax, haiku, this one, etc). I have no idea which one to use. This was the case when I first played with jax 5 years ago, and it’s still the case today (even worse it seems). This makes it a non starter for me.

Why? Use any. I used to work with Flax, now I work with Equinox more. Choose any between Flax, Equinox, and Haiku.

Use any

Too much overhead when I just want to get shit done.

Re: Jax and Equinox: What are they and why should I bother?

#15

Post by mccoyb »

spearman wrote: post #13
mccoyb wrote: post #10

Earlier quoted context omitted.

I think they’re likely using MLX in house now, no? (Probably not everyone, ofc - but seems likely that many will just use the native array framework designed explicitly for MX chips)

From https://machinelearning.apple.com/research/introducing-apple... > Our foundation models are trained on Apple's AXLearn framework, an open-source project we released in 2023. It builds on top of JAX

Wow! thanks for the ref

Re: Jax and Equinox: What are they and why should I bother?

#16

Post by __rito__ »

p1esk wrote: post #14
__rito__ wrote: post #12

Earlier quoted context omitted.

Why? Use any. I used to work with Flax, now I work with Equinox more. Choose any between Flax, Equinox, and Haiku.

Use any Too much overhead when I just want to get shit done.

Okay, I am telling you to use Flax.

Its high level API is quite similar to that of PyTorch. So, you will feel right at home.

Re: Jax and Equinox: What are they and why should I bother?

#17

Post by p1esk »

__rito__ wrote: post #16
p1esk wrote: post #14

Earlier quoted context omitted.

Use any Too much overhead when I just want to get shit done.

Okay, I am telling you to use Flax. Its high level API is quite similar to that of PyTorch. So, you will feel right at home.

So, “do as I say, not as I do”? I just looked at Equinox and it does look a little better than flax, but flax seems to be more widely used. Both haiku and trax seem to be on the way out. All four were intended to do exactly the same thing. What a mess of an ecosystem.
Post reply on HN