683 dependents
| Package | Description | Downloads/month |
|---|---|---|
| Composable transformations of Python+NumPy programs: differentiate, vectorize, J... | 20M | |
| 4.4M | ||
| Optax is a gradient processing and optimization library for JAX. | 4.1M | |
| Chex: Testing made fun, in JAX! | 2.2M | |
| Orbax provides common checkpointing and persistence utilities for JAX users | 895K | |
| CLU lets you write beautiful training loops in JAX. | 830K | |
| Accurate Quantized Training library. | 581K | |
| Probabilistic programming with NumPy powered by JAX for autograd and JIT compila... | 577K | |
| Hardware accelerated, batchable and differentiable optimizers in JAX. | 554K | |
| a Jax quantization library | 494K | |
| Task-based datasets, preprocessing, and evaluation for sequence models. | 347K | |
| jax library for E3 Equivariant Neural Networks | 237K | |
| 228K | ||
| A Pallas Custom Kernel Library. | 227K | |
| Task-based datasets, preprocessing, and evaluation for sequence models. | 216K | |
| Distrax: Probability distributions in JAX. | 206K | |
| DrJAX - Scalable and Differentiable MapReduce Primitives in JAX. | 160K | |
| a KLU solver for JAX | 149K | |
| Multi-Joint dynamics with Contact. A general purpose physics simulator. | 138K | |
| BlackJAX is a Bayesian Inference library designed for ease of use, speed and mod... | 135K | |
| Differentiable, Hardware Accelerated, Molecular Dynamics | 95K | |
| Pytrees + dataclasses ❤️ | 76K | |
| Rax is a Learning-to-Rank library written in JAX. | 65K | |
| Massively parallel rigidbody physics simulation on accelerator hardware. | 53K | |
| A library of reinforcement learning building blocks in JAX. | 37K | |
| An open-source framework for machine learning and other computations on decentra... | 33K | |
| Saddle point optimization for molecular systems | 30K | |
| Simons Stellarator Optimizer Code | 29K | |
| Dopamine is a research framework for fast prototyping of reinforcement learning ... | 26K | |
| A Python package for probabilistic state space modeling with JAX | 23K | |
| Probabilistic Programming and Nested sampling in JAX | 23K | |
| Accelerated, Python-only, single-cell integration benchmarking metrics | 22K | |
| Qrisp - The next generation of quantum algorithm development | 19K | |
| A differentiable physics engine and multibody dynamics library for control and r... | 18K | |
| easydel jax kernels writen in triton for gpus and pallas for tpus | 17K | |
| JAX-accelerated meta-reinforcement learning environments inspired by XLand and M... | 16K | |
| Sequential Least Squares Programming (SLSQP) optimizer implemented in pure JAX | 16K | |
| Differentiable QP solver in JAX | 15K | |
| PyHGF: A neural network library for predictive coding | 14K | |
| Gaussian processes in JAX and Equinox. | 13K | |
| Monte Carlo tree search in JAX | 13K | |
| 13K | ||
| A JIT compiler for hybrid quantum programs in PennyLane | 12K | |
| A patching tool to remove the numpy<2 constraint from official mediapipe wheels. | 12K | |
| GLM-HMM and GLM-HMMT tooling for behavioural task analysis. | 12K | |
| Classically postprocess noisy quantum samples to yield more accurate energy esti... | 10K | |
| The jaxnnls package provides functions to solve non-negative | 10K | |
| Accelerate, Optimize performance with streamlined training and serving options w... | 9K | |
| The Theory of Functional Connections: A functional interpolation method with app... | 8K | |
| Implementation of popular vision models in Jax | 8K |
