flaim is a library of state-of-the-art pre-trained vision models, plus common deep learning modules in computer vision, for JAX/Flax. It exposes a host of diverse image models through a straightforward interface with an emphasis on simplicity, leanness, and readability, and supplies lower-level modules for designing custom architectures.

Questions and feedback are welcome in the comments sections.