Step size adaptation for the No-U-Turn Sampler

Kolmogorov–Arnold Networks built on JAX

Simulation based calibration and generation of synthetic data.

Python wrapper for nuts-rs

ML Collections is a library of Python Collections designed for ML use cases.

Home for "How To Scale Your Model", a short blog-style textbook about scaling LLMs on TPUs

A pedagogical implementation of Autograd

High accuracy RAG for answering questions from scientific documents with citations

NeuralGCM's differentiable dycore

Probabilistic Programming and Nested sampling in JAX

pasta marker project

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

Bayesian Neural Field models for prediction in large-scale spatiotemporal datasets

Normalizing-flow enhanced sampling package for probabilistic inference in Jax

Hardware accelerated, batchable and differentiable optimizers in JAX.

tree is a library for working with nested data structures

Fast Bayesian Inference in Python. 50+ conjugate models with vectorized updates, sufficient statistic helpers, built-in plotting, and SciPy integration.

State of the art inference for your bayesian models.

Application of the L2HMC algorithm to simulations in lattice QCD.

Optax is a gradient processing and optimization library for JAX.

Inference Combinators in JAX

Python-based research interface for blackbox and hyperparameter optimization, based on the internal Google Vizier Service.

an ambient intelligence library

Optimal transport tools implemented with the JAX framework, to solve large scale matching problems of any flavor.

LLM(😽)