A simple library for scaling up JAX programs

selfEEG: a Python library for Self-Supervised Learning on Electroencephalography (EEG) data

Code repository for Particle Denoising Diffusion Sampler

diffusionjax is a simple and accessible diffusion models package in JAX

Very deep VAEs in JAX/Flax

Differentiable, Hardware Accelerated, Molecular Dynamics

Graph Learning with JAX

Implementation of Action Matching for the Schrödinger equation

Implementation of Action Matching

Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax

Initial public release of code, data, and model weights for FourCastNet

A family of diffusion models for text-to-audio generation.

A latent text-to-image diffusion model

🎢 Creating and sharing simulation environments for embodied and synthetic data research

Learning the optimal transport map via input convex neural neworks

Euclidean Wasserstein-2 optimal transportation

Repository containing Google Colab code for "Debiaser Beware: Pitfalls of Centering Regularized Transport Maps", which was recently accepted to ICML 2022.

Implementation of the 1-Nearest-Neighbor estimator for estimating optimal transport maps, as seen in [Plugin Estimation of Smooth Transport Maps](https://arxiv.org/abs/2107.12364) (Manole et al.; 2…

Pytorch implementation of diffusion models on Lie Groups for 6D grasp pose generation https://sites.google.com/view/se3dif/home

Audio generation using diffusion models, in PyTorch.

TorchGeo: datasets, samplers, transforms, and pre-trained models for geospatial data

Transformers with doubly stochastic attention

PyTorch Implementation of DSB for Score Based Generative Modeling. Experiments managed using Hydra.

Generative Modeling with Optimal Transport Maps - ICLR 2022

Tackling the Generative Learning Trilemma with Denoising Diffusion GANs https://arxiv.org/abs/2112.07804

Score-Based Generative Modeling with Critically-Damped Langevin Diffusion