Pinned Loading

1. psgd_jax psgd_jax Public

   Implementation of PSGD optimizer in JAX

   Python 19 1

2. kron_torch kron_torch Public

   An implementation of PSGD Kron second-order optimizer for PyTorch

   Python 16 2

3. llm-jax llm-jax Public

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

   Python 9

4. image-classification-jax image-classification-jax Public

   Image classification in JAX with ViT, resnet, cifar10, cifar100, imagenette, and imagenet

   Python 2