Deep Learning for humans

🏛️A research-friendly codebase for fast experimentation of single-agent reinforcement learning in JAX • End-to-End JAX RL

The Unified ML Representation

Deep Gaussian Markov Random Fields and their extensions

Optimal transport tools implemented with the JAX framework, to get differentiable, parallel and jit-able computations.

A simple JAX-based implementation of random search for locomotion tasks using MuJoCo XLA (MJX).

🤗 Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch and FLAX.

🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.

An open-source JAX-based gravitational-wave population inference toolkit

Deep Learning examples using the Jax ecosystem of libraries

JAX compilation of RDDL description files, and a differentiable planner in JAX.

Implementation of DETR in JAX

Accelerate your training with this open-source library. Optimize performance with streamlined training and serving options with JAX. 🚀

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

A library for scientific machine learning and physics-informed learning

my jax fork

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

A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper and Ada GPUs, to provide better performance with lower memory utilization in both training and inference.

A simple, easy-to-understand library for diffusion models using Flax and Jax. Includes detailed notebooks on DDPM, DDIM, and EDM with simplified mathematical explanations. Made as part of my journey for learning and experimenting with state-of-the-art generative AI.

NAACL '24 (Demo) / MlSys @ NeurIPS '23 - RedCoast: A Lightweight Tool to Automate Distributed Training and Inference