Plugging Into JAX by Nick Doiron

medium.comsite

Compares Flax, Haiku, and Objax on the Kaggle flower classification challenge.

from

JAX

added

2026-10-10

likes

0

similar

1. ▲
   Flax NNX github.com

   site category #jax link ↗

   An evolution on Flax by the same team

2. ▲
   Elegy poets-ai.github.io

   site category #jax link ↗

   A High Level API for Deep Learning in JAX. Supports Flax, Haiku, and Optax.

3. ▲
   FLAX github.com

   tool category #python-data-science link ↗

   A neural network library for JAX that is designed for flexibility.

4. ▲
   Lorax github.com

   tool category #jax link ↗

   Automatically apply LoRA to JAX models (Flax, Haiku, etc.)

   
5. ▲
   Flax Models github.com

   site category #jax link ↗

   Collection of models and methods implemented in Flax.

6. ▲
   flaxmodels github.com

   tool category #jax link ↗

   Pretrained models for Jax/Flax.

JAX › Tutorials and Blog Posts: “Compares Flax, Haiku, and Objax on the Kaggle flower classification challenge.”