JAX

exojaxHajimeKawahara/exojax

A JAX-based differentiable spectral modeling library for exoplanets, brown dwarfs, and M dwarfs.

KernexASEM000/kernex

Kernex extends JAX with kmap and kscan for differentiable stencil computations, enabling efficient array transformations.

GNNs for Solving Combinatorial Optimization ProblemsIvanIsCoding/GNN-for-Combinatorial-Optimization

A JAX + Flax implementation of physics-inspired graph neural networks for solving combinatorial optimization problems like Max-Cut and Maximum Independent Set.

delta PVromanodev/deltapv

A photovoltaic simulator with automatic differentiation for solar cell modeling and optimization, built on JAX.

astronomixleo1200/astronomix

A differentiable hydrodynamics and magnetohydrodynamics code for astrophysics built with JAX, enabling gradient-based inverse modeling and multi-GPU simulations.

PGMaxvicariousinc/PGMax

A JAX-based library for loopy belief propagation on discrete factor graphs, enabling efficient probabilistic inference.

DiffeRTjeertmans/DiffeRT

A differentiable ray tracing toolbox for radio propagation simulations, built on JAX for optimization and machine learning.

JAXFitdipolar-quantum-gases/jaxfit

GPU/TPU accelerated nonlinear least-squares curve fitting using JAX, designed as a drop-in replacement for SciPy's curve_fit.

sklearn-jax-kernelsExpectationMax/sklearn-jax-kernels

Composable kernels for scikit-learn implemented in JAX, enabling faster kernel computations and automatic differentiation.

FlaxVisionrolandgvc/flaxvision

A collection of neural network models ported from torchvision for use with JAX and Flax.

imax4rtemi5/imax

A JAX-based library for fast, composable image augmentation with geometric and color transformations.

Coreaxgchq/coreax

A JAX-based library for coreset algorithms that reduce large datasets to smaller representative subsets while preserving statistical properties.

vivsimhaimingz/vivsim

A Python library for accelerated fluid-structure interaction simulations using the immersed boundary lattice Boltzmann method, powered by JAX.

flaxdiffAshishKumar4/FlaxDiff

A simple, easy-to-understand library for diffusion models built with Flax and JAX, focusing on readability and learning.

tmmaxbahremsd/tmmax

A high-performance JAX-based library for computing optical properties of multilayer thin-film structures using the transfer matrix method.

JAX-in-Celluwplasma/JAX-in-Cell

A 1D3V particle-in-cell simulation code for plasma physics, accelerated using JAX for performance on CPUs, GPUs, and TPUs.

FID computationmatthias-wright/jax-fid

A JAX/Flax implementation of the Fréchet Inception Distance (FID) metric for evaluating generative models.

DeepSeek-R1-Flax-1.5B-DistillJ-Rosser-UK/Torch2Jax-DeepSeek-R1-Distill-Qwen-1.5B

Flax (JAX) and PyTorch implementations of the DeepSeek-R1-Distill-Qwen-1.5B language model with weight conversion utilities.

Tutorial: image classification with JAX and Flax Linen by 8bitmp38bitmp3/JAX-Flax-Tutorial-Image-Classification-with-Linen

A tutorial on building and training a convolutional neural network for MNIST image classification using Flax Linen and Optax.

brainunitchaobrain/brainunit

A physical units and unit-aware mathematical system built on JAX for brain dynamics modeling and general scientific computing.

dendritexchaobrain/dendritex

A unified interface for modeling single- and multi-compartment Hodgkin-Huxley neuron models in JAX.

DETRMasterSkepticista/detr

A minimal JAX/Flax implementation of DETR with optimizations like Flash Attention and Sinkhorn solver.

braintaichichaobrain/braintaichi

A Python library that uses Taichi Lang to create custom, high-performance operators for brain dynamics simulations.

foragaxi-m-iron-man/Foragax

An Agent-Based Modelling framework built on JAX for scalable and efficient simulations with automatic vectorization and JIT compilation.

jax-raftalebeck/jax-raft

A JAX/Flax implementation of the RAFT optical flow estimator with ported checkpoints and reproducible results.

Two Player Auction Learningdegregat/two-player-auctions

JAX/Haiku implementation of a two-player game approach to learning optimal multi-bidder, multi-item auctions.