Evan Walters

kron_torch. An implementation of PSGD Kron second-order optimizer for PyTorch

psgd_jax. Implementation of PSGD optimizer in JAX

llm-jax. Train a SmolLM-style llm on fineweb-edu in JAX/Flax with an assortment of optimizers.

AdaMod-tf. Tensorflow version of the AdaMod Optimizer

distributed_kron. An implementation of PSGD Kron in JAX for distributed training in JAX or Flax

quad_torch. An implementation of PSGD-QUAD optimizer for PyTorch

diffGrad-tf. DiffGrad Optimizer in Tensorflow

jax-torch-comparison. Python

sophia-schedulefree-jax. Sophia optimizer with ScheduleFree

xLSTM-JAX. An implementation of mLSTM from xLSTM in JAX

image-classification-jax. Image classification in JAX with ViT, resnet, cifar10, cifar100, imagenette, and imagenet

Apollo-tf. Apollo optimizer in tensorflow

tiny-muon-jax. Python

flat-sophia. sophia optimizer further projected towards flat areas of loss landscape

easydel-experiments.

Rectified-diffGrad-tf. Rectified Adam optimizer with diffGrad Optimizer modification

ordax. A collection of second-order optimizers and experiments in JAX

flash-muon. Flash-Muon: An Efficient Implementation of Muon Optimizer

torch-grokking. Grokking on modular arithmetic in less than 150 epochs in MLX

gidd. Code accompanying the paper "Generalized Interpolating Discrete Diffusion"

quad_jax. An implementation of PSGD-QUAD optimizer for JAX

big_vision. Official codebase used to develop Vision Transformer, SigLIP, MLP-Mixer, LiT and more.

EasyLM. Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.

jax_sourceror. Turn jitted jax functions back into python source code

clip-jax. Train vision models using JAX and 🤗 transformers

I-xTx-jax. playing around with minimizing the toy fn I - x ^ T x

gpt-jax. Jax/Flax rewrite of Karpathy's nanoGPT

dit_flow. DiT (training + flow matching) in Jax

levanter. Legible, Scalable, Reproducible Foundation Models with Named Tensors and Jax