JAX
216 items from n2cholas/awesome-jax ★2,167
-
Lorax github.com
Automatically apply LoRA to JAX models (Flax, Haiku, etc.)
-
Spyx github.com
Spiking Neural Networks in JAX for machine learning on neuromorphic hardware.
-
JaxDF github.com
Framework for differentiable simulators with arbitrary discretizations.
-
SymJAX github.com
Symbolic CPU/GPU/TPU programming.
-
jax-tqdm github.com
Add a tqdm progress bar to JAX scans and loops.
-
TF2JAX github.com
Convert functions/graphs to JAX functions.
-
jax-resnet github.com
Implementations and checkpoints for ResNet variants in Flax.
-
Eqxvision github.com
Equinox version of Torchvision.
-
Einshape github.com
DSL-based reshaping library for JAX and other frameworks.
-
bayex github.com
Bayesian Optimization powered by JAX.
-
jax-unirep github.com
Library implementing the UniRep model for protein machine learning applications.
-
CR.Sparse github.com
XLA accelerated algorithms for sparse representations and compressive sensing.
-
kalman-jax github.com
Approximate inference for Markov (i.e., temporal) Gaussian processes using iterated Kalman filtering and smoothing.
-
XMC-GAN github.com
Official implementation of Cross-Modal Contrastive Learning for Text-to-Image Generation.
-
NuX github.com
Normalizing flows with JAX.
-
efax github.com
Exponential Families in JAX.
-
exojax github.com
Automatic differentiable spectrum modeling of exoplanets/brown dwarfs compatible to JAX.
-
Kernex github.com
Differentiable stencil decorators in JAX.
-
astronomix github.com
differentiable (magneto)hydrodynamics for astrophysics in JAX
-
delta PV github.com
A photovoltaic simulator with automatic differentation.
-
GNNs for Solving Combinatorial Optimization Problems github.com
A JAX + Flax implementation of Combinatorial Optimization with Physics-Inspired Graph Neural Networks.
-
PGMax github.com
A framework for building discrete Probabilistic Graphical Models (PGM's) and running inference inference on them via JAX.
-
DiffeRT github.com
Differentiable Ray Tracing toolbox for Radio Propagation powered by the JAX ecosystem.
-
JAXFit github.com
Accelerated curve fitting library for nonlinear least-squares problems (see arXiv paper).
-
sklearn-jax-kernels github.com
scikit-learn kernel matrices using JAX.
-
safejax github.com
Serialize JAX, Flax, Haiku, or Objax model params with 🤗safetensors.
-
FlaxVision github.com
Flax version of TorchVision.
-
vivsim github.com
Fluid-structure interaction simulations using Immersed Boundary-Lattice Boltzmann Method.
-
imax github.com
Image augmentations and transformations.
-
Coreax github.com
Algorithms for finding coresets to compress large datasets while retaining their statistical properties.
-
tmmax github.com
Vectorized calculation of optical properties in thin-film structures using JAX. Swiss Army knife tool for thin-film optics research
-
flaxdiff github.com
Framework and Library for building and training Diffusion models in multi-node multi-device distributed settings (TPUs)
-
JAX-in-Cell github.com
Plasma physics simulations using a PIC (Particle-in-Cell) method to self-consistently solve for electron and ion dynamics in electromagnetic fields
-
lqg github.com
Official implementation of Bayesian inverse optimal control for linear-quadratic Gaussian problems from the paper Putting perception into action with inverse optimal control for continuous psychophysics
-
FID computation github.com
Port of mseitzer/pytorch-fid to Flax.
-
DeepSeek-R1-Flax-1.5B-Distill github.com
Flax implementation of DeepSeek-R1 1.5B distilled reasoning LLM.
-
Tutorial: image classification with JAX and Flax Linen by 8bitmp3 github.com
Learn how to create a simple convolutional network with the Linen API by Flax and train it to recognize handwritten digits.
-
MBIRJAX github.com
High-performance tomographic reconstruction.
-
brainstate github.com
State-based Transformation System for Program Compilation and Augmentation.
-
brainunit github.com
Physical units and unit-aware mathematical system in JAX.
-
dendritex github.com
Dendritic Modeling in JAX.
-
DETR github.com
Flax implementation of DETR: End-to-end Object Detection with Transformers using Sinkhorn solver and parallel bipartite matching.
-
braintaichi github.com
Leveraging Taichi Lang to customize brain dynamics operators.
-
foragax github.com
Agent-Based modelling framework in JAX.
-
jax-raft github.com
JAX/Flax port of the RAFT optical flow estimator.
-
JaxNeuS github.com
Implementation of NeuS: Learning Neural Implicit Surfaces by Volume Rendering for Multi-view Reconstruction
-
MLP Mixer github.com
Minimal implementation of MLP-Mixer: An all-MLP Architecture for Vision.
-
gMLP github.com
Implementation of Pay Attention to MLPs.
-
jax-flows github.com
Normalizing flows in JAX.
-
Two Player Auction Learning github.com
JAX implementation of the paper Auction learning as a two-player game.
