This repository is a curated list of resources for JAX, a high-performance numerical computation library for machine learning research. It serves as a comprehensive catalog for developers and researchers looking to leverage JAX's automatic differentiation and XLA compilation capabilities on accelerators like GPUs and TPUs.

How It Works

The list is organized into categories such as Libraries, Models and Projects, Videos, Papers, and Tutorials. It highlights various JAX-based libraries for neural networks (Flax, Haiku, Equinox), probabilistic programming (NumPyro), reinforcement learning (RLax, Coax, gymnax), and scientific computing (JAX, M.D., SCICO). The curated nature aims to provide a structured overview of the rapidly growing JAX ecosystem.

Quick Start & Requirements

• Installation: JAX itself is typically installed via pip: pip install jax[cuda12_pip] (for CUDA 12). Specific libraries may have their own installation instructions.
• Prerequisites: Python 3.8+, NumPy. GPU support requires CUDA and cuDNN. TPU support is available via Google Cloud.
• Resources: Setup time varies based on project complexity. Many projects include Colab notebooks for quick experimentation.
• Links: JAX GitHub, Awesome JAX

Highlighted Details

• Extensive coverage of neural network libraries like Flax, Haiku, and Equinox.
• Numerous implementations of cutting-edge research papers in computer vision, NLP, and physics.
• Rich collection of video tutorials, papers, and blog posts for learning and advanced usage.
• Includes specialized libraries for graph neural networks, probabilistic programming, and reinforcement learning.

Maintenance & Community

• Actively maintained and community-driven, with contributions welcome.
• Links to community resources like GitHub Discussions and an unofficial JaxLLM Discord are provided.

Licensing & Compatibility

• JAX is typically released under the Apache 2.0 license.
• Individual libraries listed may have different licenses; users should verify compatibility for commercial or closed-source use.

Limitations & Caveats

This is a curated list, not a single library. Users must evaluate each listed project individually for its maturity, maintenance status, and specific requirements.