mattjj

autodidact

A pedagogical implementation of Autograd

pyhsmm

svae

code for Structured Variational Autoencoders

pybasicbayes

pyslds

pylds

some tools for gaussian linear dynamical systems

pyhsmm-autoregressive

autoregressive plugin

ode-diff-notes

pplham

spectral_clustering

oryx

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

kalman_grads

autograd_linalg

improve scipy.linalg and corresponding autograd gradfuns

computation-thru-dynamics

Understanding computation in artificial and biological recurrent networks through the lens of dynamical systems.

axlearn

An Extensible Deep Learning Library

dex-lang

Research language for array processing in the Haskell/ML family

jax

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

jax-llm-examples

pypolyagamma

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

flax

Flax is a neural network library for JAX that is designed for flexibility.

jax-md

Differentiable, Hardware Accelerated, Molecular Dynamics

quax

Multiple dispatch over abstract array types in JAX.

scaling-book

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

blog-benchmark-rwmetropolis

Code to reproduce the results of https://rlouf.github.io/post/jax-random-walk-metropolis/

flash-attention-jax

Implementation of Flash Attention in Jax

HighPerfLLMs2024

iree

👻

jax-triton

jax-triton contains integrations between JAX and OpenAI Triton

julia

The Julia Language: A fresh approach to technical computing.

mattjj.github.io