JAX

jaxliebrentyi/jaxlie

A JAX library implementing Lie groups for rigid body transformations in computer vision and robotics.

mcxrlouf/mcx

A JAX-powered probabilistic programming library focused on performant sampling methods for Bayesian inference on CPU, GPU, and TPU.

KFAC-JAXdeepmind/kfac-jax

A JAX library for second-order optimization of neural networks using the K-FAC curvature approximation algorithm.

FDTDXymahlau/fdtdx

An efficient open-source Python package for 3D photonic nanostructure simulation and design using GPU-accelerated FDTD with automatic differentiation.

dynamiqsdynamiqs/dynamiqs

A Python library for GPU-accelerated and differentiable quantum systems simulation built with JAX.

Equivariant MLPmfinzi/equivariant-MLP

A JAX library for automatically generating equivariant neural network layers for arbitrary symmetry groups via constraint solving.

FedJAXgoogle/fedjax

A JAX-based library for federated learning simulations that emphasizes ease-of-use in research.

flaxmodelsmatthias-wright/flaxmodels

A collection of pretrained deep learning models (StyleGAN2, GPT2, VGG, ResNet) for the Jax/Flax ecosystem.

jax-cosmoDifferentiableUniverseInitiative/jax_cosmo

A differentiable cosmology library built with JAX for automatic differentiation of cosmological calculations.

jaxnsJoshuaalbert/jaxns

A JAX-based probabilistic programming framework using nested sampling for fast Bayesian inference and evidence computation.

flowjaxdanielward27/flowjax

A JAX library for distributions, bijections, and normalizing flows implemented as Equinox modules.

torchaxgoogle/torchax

A PyTorch frontend for JAX that enables running PyTorch code on TPUs and provides seamless PyTorch-JAX interoperability.

Optimal Transport Toolsgoogle-research/ott

A Python toolbox for solving optimal transport problems with JAX-powered computational efficiency.

tree-mathgoogle/tree-math

A Python library that enables mathematical operations on JAX pytrees, allowing numerical algorithms to work directly with complex nested data structures.

jwaveucl-bug/jwave

A JAX-based research framework for differentiable and parallelizable acoustic simulations, running on CPU, GPU, and TPU.

NesTgoogle-research/nested-transformer

A vision transformer architecture that aggregates nested local transformers on image blocks for better accuracy, data efficiency, and convergence.

Praxisgoogle/praxis

A layer library for JAX-based machine learning projects, optimized for large-scale ML.

NAVIXepignatelli/navix

A JAX-powered reimplementation of MiniGrid offering over 1000x speedup for reinforcement learning experiments.

kvaxnebius/kvax

A FlashAttention 2 implementation for JAX with block-wise document mask optimization and context parallelism for efficient long-sequence training.

SCICOlanl/scico

A Python package built on JAX for solving inverse problems in scientific imaging using optimization and prior models.

jax-modelsDarshanDeshpande/jax-models

Unofficial JAX/Flax implementations of deep learning research papers for vision transformers and other architectures.

Parallaxsrush/parallax

A pure, immutable module system for JAX that replaces PyTorch-style imperative coding with declarative parameter trees.

Loraxdavisyoshida/lorax

A JAX transform that implements LoRA (Low-Rank Adaptation) for efficient fine-tuning of large models with minimal memory overhead.

Spyxkmheckel/spyx

A compact spiking neural network library built on JAX and Haiku, offering high-performance training via surrogate gradient descent and neuroevolution.

JaxDFucl-bug/jaxdf

A JAX-based framework for building differentiable numerical simulators with arbitrary discretizations for physical systems.

SymJAXSymJAX/SymJAX

A symbolic programming library built on JAX for concise, explicit, and optimized machine learning computations.

jax-tqdmjeremiecoullon/jax-tqdm

Add tqdm progress bars to JAX scans and loops using decorators, enabling side-effect-free progress tracking.

TF2JAXdeepmind/tf2jax

An experimental library that converts TensorFlow functions and graphs into JAX functions for reuse and fine-tuning within JAX codebases.

jax-resnetn2cholas/jax-resnet

Flax implementations and pretrained checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX.

Eqxvisionpaganpasta/eqxvision

A Python package providing popular computer vision model architectures built with Equinox for JAX.

bayexalonfnt/bayex

A lightweight Bayesian optimization library built on JAX for efficient optimization of expensive-to-evaluate functions.

jax-unirepElArkk/jax-unirep

A performant JAX reimplementation of the UniRep model for generating protein sequence representations.

kalman-jaxAaltoML/kalman-jax

A JAX-based framework for approximate inference in Markov Gaussian processes using iterated Kalman smoothing.

XMC-GANgoogle-research/xmcgan_image_generation

Official JAX implementation of XMC-GAN for text-to-image generation using cross-modal contrastive learning.

NuXInformation-Fusion-Lab-Umass/NuX

A JAX library for building, training, and evaluating normalizing flows for probabilistic modeling.

Rockpoolsynsense/rockpool

A Python library for building, training, and deploying spiking neural networks with support for multiple simulation backends and neuromorphic hardware.