Pinned Loading

1. distributed_kron distributed_kron Public

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

   Python 1

2. psgd_jax psgd_jax Public

   Implementation of PSGD optimizer in JAX

   Python 26 2

3. kron_torch kron_torch Public

   An implementation of PSGD Kron second-order optimizer for PyTorch

   Python 21 2

4. llm-jax llm-jax Public

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

   Python 14

5. image-classification-jax image-classification-jax Public

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

   Python 2