jax

  • 17th July 2026

fae: FLUX.1 SAE training on TPUs

jax diffusion wip

An implementation of MMDiT in jax, + SAE training on TPUs

Read more →
  • 5th August 2025

saex: training SAEs in Jax

interpretability jax sae

A library for training Sparse Autoencoders in Jax

Read more →
  • 3rd August 2025

jif: Jax text diffusion

jax diffusion

Toy implementation of masked diffusion language models

Read more →