Mathematical operations for JAX pytrees

A machine learning compiler for GPUs, CPUs, and ML accelerators

Gaussian processes in JAX and Flax.

Flax is a neural network library for JAX that is designed for flexibility.

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

MLCommons Algorithmic Efficiency is a benchmark and competition measuring neural network training speedups due to algorithmic improvements in both training algorithms and models.

A playbook for systematically maximizing the performance of deep learning models.

Hardware accelerated, batchable and differentiable optimizers in JAX.

Rax is a Learning-to-Rank library written in JAX.

Flax is a neural network library for JAX that is designed for flexibility.

Research language for array processing in the Haskell/ML family

Concise deep learning for JAX

Robust Bi-Tempered Logistic Loss Based on Bregman Divergences. https://arxiv.org/pdf/1906.03361.pdf

Differentiable, Hardware Accelerated, Molecular Dynamics

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more