285 dependents
| Package | Description | Downloads/month |
|---|---|---|
| Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax | 272K | |
| Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox. https:... | 208K | |
| Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capab... | 180K | |
| [ICML'26] Phonon fine-tuning (PFT) and [NeurIPS'25 AI4Mat] Nequix: Training a fo... | 85K | |
| Named Tensors for Legible Deep Learning in JAX | 75K | |
| Interpolation and function approximation with JAX | 56K | |
| export JAX to ONNX | 37K | |
| Legible, Scalable, Reproducible Foundation Models with Named Tensors and Jax | 34K | |
| Numerical quadrature with JAX | 29K | |
| Orthogonal polynomials with JAX | 22K | |
| Differentiable Ray Tracing Toolbox for Radio Propagation Simulations | 17K | |
| Sequential Least Squares Programming (SLSQP) optimizer implemented in pure JAX | 16K | |
| High-performance quantum systems simulation with JAX (GPU-accelerated & differen... | 14K | |
| Gaussian processes in JAX and Equinox. | 13K | |
| Efficient Differentiable n-d PDE Solvers in JAX. | 13K | |
| Machine learning algorithms for many-body quantum systems | 13K | |
| Neural Emulator Architectures in JAX. | 13K | |
| [Neurips 2024] A benchmark suite for autoregressive neural emulation of PDEs. (≥... | 13K | |
| Neural Emulator Architectures in JAX. | 13K | |
| Stellarator Equilibrium and Optimization Suite | 12K | |
| Normalizing-flow enhanced sampling package for probabilistic inference in JAX | 10K | |
| Fourier interpolation and function approximation with JAX | 9K | |
| Multiple dispatch over abstract array types in JAX. | 9K | |
| A small library of paramaterizations and parameter constraints for PyTrees. | 8K | |
| Implementation of popular vision models in Jax | 8K | |
| Unitful Quantities in JAX | 8K | |
| Fast universal quantum circuit sampler for QEC based on ZX stabilizer rank decom... | 8K | |
| A Python implementation of active inference for Markov Decision Processes | 8K | |
| The tiniest of Gaussian Process libraries | 7K | |
| Parametric modeling in JAX + Equinox | 7K | |
| Easy to use distributions, bijections and normalizing flows in JAX. | 7K | |
| Pre-quaxed libraries for use in multiple dispatch over abstract array types in J... | 7K | |
| JAX backend for QuTiP | 5K | |
| Bayesian Inference with JAX | 5K | |
| Parametric radio frequency modeling, optimization and sampling | 4K | |
| JAX library for optimization and export of models for use with the UZU inference... | 4K | |
| Python wrapper for yield input packages | 4K | |
| A high-performance Python package for estimating latent-class conditional logit ... | 3K | |
| Coordinates in JAX | 3K | |
| Physics Informed Neural Network with JAX | 3K | |
| Collimator.ai core simulation engine and API client | 3K | |
| JAX-native differentiable B-Rep kernel | 3K | |
| Python+JAX package for differentiable CMB+BBN computations. See https://abcmb.re... | 3K | |
| A library for fast Jax experimentation | 3K | |
| Gyrokinetic Spectral Solver of Plasma Turbulence | 3K | |
| Deep equilibrium models in JAX/Equinox | 2K | |
| JAX-powered brain network simulation with gradient-based optimization | 2K | |
| Turn SymPy expressions into trainable JAX expressions. | 2K | |
| differentiable (magneto)hydrodynamics for astrophysics in JAX | 2K | |
| A Jaxley-based library of ion channels and syanpses for biophysical neuron model... | 2K |
