JaxNN: Foundation Models in JAX/Flax

Jax Neural Networks is an open-source library for foundation models in JAX and Flax. It provides a unified framework for loading, creating, and using pretrained models (e.g., ResNet, ViT).