Elegy

poets-ai.github.iosite

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

from

JAX

added

2026-10-10

likes

0

similar

1. ▲
   Plugging Into JAX by Nick Doiron medium.com

   site category #jax link ↗

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

2. ▲
   FedJAX github.com

   tool category #jax link ↗

   Federated learning in JAX, built on Optax and Haiku.

   
3. ▲
   FLAX github.com

   tool category #python-data-science link ↗

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

4. ▲
   Using JAX to accelerate our research by David Budden and Matteo Hessel deepmind.com

   site category #jax link ↗

   Describes the state of JAX and the JAX ecosystem at DeepMind.

   
5. ▲
   Deep Learning tutorials with JAX+Flax by Phillip Lippe uvadlc-notebooks.readthedocs.io

   site category #jax link ↗

   A series of notebooks explaining various deep learning concepts, from basics (e.g. intro to JAX/Flax, activiation functions) to recent advances (e.g., Vision Transformers, SimCLR), with translations to PyTorch.

6. ▲
   JAX, Flax & Transformers 🤗 github.com

   site category #jax link ↗

   3 days of talks around JAX / Flax, Transformers, large-scale language modeling and other great topics.

JAX › Libraries: “A High Level API for Deep Learning in JAX. Supports Flax, Haiku, and Optax.”