An implementation of MuZero in JAX.

🦁 A research-friendly codebase for fast experimentation of multi-agent reinforcement learning in JAX

Quantile regression-based reinforcement learning algorithms for optimizing different risk measures

Code for the paper: "Generalized Discrete Diffusion from Snapshots"

Code for TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion

[ICML2025] Official implementation of Efficient Online Reinforcement Learning for Diffusion Policies appearing in ICML 2025.

Self-contained, minimalistic implementation of diffusion models with Pytorch.

Simple Guidance Mechanisms for Discrete Diffusion Models

An implementation of a small scale autoencoder and latent diffusion model

Diffusion Probabilistic Model in Jax

Score-based Diffusion models in JAX.

Generation of protein sequences and evolutionary alignments via discrete diffusion models

Official implementation of the paper: "DeFoG: Discrete Flow Matching for Graph Generation"

A minimal PyTorch re-implementation of AlphaFold2's model & training

Official Implementation of iMF https://arxiv.org/abs/2512.02012

Official Implementation of pMF https://arxiv.org/abs/2601.22158

Implementation of One-step Latent-free Image Generation with Pixel Mean Flows (Lu et al., 2026)

code for "Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching"

Adjoint Matching Colab

Q-learning with Adjoint Matching

[ICML 2024 Best Paper] Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution (https://arxiv.org/abs/2310.16834)

Educational implementation of the Discrete Flow Matching paper

Single-step image generation at 306 FPS. Drifting vs Diffusion head-to-head on CIFAR-10.

Code for the paper https://arxiv.org/abs/2402.04997

Generative Modeling via Drifting in MLX

Official PyTorch implementation of Vision DiffMask, a post-hoc interpretation method for vision models.