GitHub repository: phlippe/jax_trainer
Primary language: Python
Description: Lightning-like training API for JAX with Flax