[HN Gopher] Jax and Equinox: What are they and why should I bother?
___________________________________________________________________
Jax and Equinox: What are they and why should I bother?
Author : spearman
Score : 32 points
Date : 2024-09-09 14:46 UTC (3 days ago)
(HTM) web link (www.garymm.org)
(TXT) w3m dump (www.garymm.org)
| atorodius wrote:
| Personally a big fan of Flax. The way it separates the params and
| the compute graph is - imo - the right (TM) way. Saying this
| after many years of soing ML :)
| ubj wrote:
| There's quite a few other libraries associated with Equinox in
| the JAX ecosystem:
|
| https://github.com/patrick-kidger/equinox?tab=readme-ov-file...
|
| I've enjoyed using Equinox and Diffrax for performing ODE
| simulations. To my knowledge the only other peer library with
| similar capabilities is the Julia DifferentialEquations.jl
| package.
| Lerc wrote:
| 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?
| mccoyb wrote:
| Equinox has great idioms -- really pioneered the Pytree
| perspective. Penzai is also great.
|
| JAX feels close to achieving a sort of high-level GPGPU
| ecosystem. Super fledgling -- but I keep finding more little
| libraries that build on JAX (and can be compositionally used with
| other JAX libraries because of it).
|
| Only problem is that lots of compositional usage leads to big
| code and therefore big compile times for XLA.
___________________________________________________________________
(page generated 2024-09-12 23:00 UTC)