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

⚡ A Fast, Extensible Progress Bar for Python and CLI

Pretrain, finetune ANY AI model of ANY size on 1 or 10,000+ GPUs with zero code changes.

Efficiently computes derivatives of NumPy code.

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.

Optax is a gradient processing and optimization library for JAX.

Spearmint Bayesian optimization codebase

Spearmint is a package to perform Bayesian optimization according to the algorithms outlined in the paper: Practical Bayesian Optimization of Machine Learning Algorithms. Jasper Snoek, Hugo Laroche…

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

A Python package for probabilistic state space modeling with JAX

Gaussian processes in JAX and Flax.

A deep learning framework for multi-animal pose tracking.

Implementation of https://srush.github.io/annotated-s4

Probabilistic Numerics in Python.

code for Structured Variational Autoencoders

A python implementation of Counterfactual Regret Minimization for poker

Accelerated pose estimation and tracking using semi-supervised learning

Arsenal of python utilities.

Python framework for inference in Hawkes processes.

Official JAX implementation of xLSTM including fast and efficient training and inference code. 7B model available at https://huggingface.co/NX-AI/xLSTM-7b.

some tools for gaussian linear dynamical systems

YASS: Yet Another Spike Sorter

Dependent multinomials made easy: stick-breaking with the Pólya-gamma augmentation

Poisson Identifiable VAE (pi-VAE)

WIP: Python port of Kilosort2

Jax SSM Library

Interpretable neural spike train models with fully-Bayesian inference algorithms