JAX-accelerated Meta-Reinforcement Learning Environments Inspired by XLand and MiniGrid 🏎️

Present a novel toy dataset constructed in embedding space for transformer testing. Also present high level qualitative results.

Large-scale text-video dataset. 10 million captioned short videos.

Bind any function written in another language to JAX with support for JVP/VJP/batching/jit compilation

XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML

[ICML 2025 Spotlight] Training Dynamics of In-Context Learning in Linear Attention.

A library for mechanistic interpretability of GPT-style language models

bluetooth mesh chat, IRC vibes

Plotting atmospheric data from the TROPOMI satellite. Plots can be imported into Unity.

Hierarchical Reasoning Model Official Release

free finite-difference time-domain (FDTD) software for electromagnetic simulations

Package for writing high-level code for parallel high-performance stencil computations that can be deployed on both GPUs and CPUs

Official implementation of the paper "Watermarking Autoregressive Image Generation" (NeurIPS'25)

Flash Attention in ~100 lines of CUDA (forward pass only)

Latent Program Network (from the "Searching Latent Program Spaces" paper)

Annotated version of the Mamba paper

Implementation of a Vision transformer in JAX NNX

Introduction to SciML with JAX

Minimal JAX implementation of k-nearest neighbors using a k-d tree.

Liquid Structural State-Space Models

Minimal yet performant LLM examples in pure JAX

Causal depthwise conv1d in CUDA, with a PyTorch interface

The Electricity Transformer dataset is collected to support the further investigation on the long sequence forecasting problem.

Weight-Space Linear Recurrent Neural Networks

LLM inference in C/C++