Oryx is a library for probabilistic programming and deep learning built on top of Jax.

A Python package of computer vision models for the Equinox ecosystem.

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

Normalizing Flows using JAX

BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.

Express & compile probabilistic programs for performant inference on CPU & GPU. Powered by JAX.

This is a collection of code samples aimed at illustrating temporal parallelization methods for sequential data.

Folklore facts on probability distribution learning, testing, and whatever-ing

Implementation of Vision Transformers in Flax

Optax is a gradient processing and optimization library for JAX.

Normalizing flows in PyTorch

Manifold-learning flows (ℳ-flows)

Code for the article "What if Neural Networks had SVDs?", to be presented as a spotlight paper at NeurIPS 2020.

JAX-based neural network library

Fast C code for sampling Polya-gamma random variates. Builds on Jesse Windle's BayesLogit library.

📐 Numerical integration (quadrature, cubature) in Python

MATLAB codes for matrix exponential times a vector.

Code for Neural Spline Flows paper

Trax — Deep Learning with Clear Code and Speed

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Must-read papers on graph neural networks (GNN)

Bayesian learning and inference for state space models