Stars
JAX-accelerated Meta-Reinforcement Learning Environments Inspired by XLand and MiniGrid 🏎️
(CVPR 2025) Switti: Designing Scale-Wise Transformers for Text-to-Image Synthesis
Loopy belief propagation for factor graphs on discrete variables in JAX
Optax is a gradient processing and optimization library for JAX.
Multi-framework implementation of Deep Kernel Shaping and Tailored Activation Transformations, which are methods that modify neural network models (and their initializations) to make them easier to…
🕹️ A diverse suite of scalable reinforcement learning environments in JAX
Second Order Optimization and Curvature Estimation with K-FAC in JAX.
Label-Efficient Semantic Segmentation with Diffusion Models (ICLR'2022)
[NeurIPS'21 Outstanding Paper] Library for reliable evaluation on RL and ML benchmarks, even with only a handful of seeds.
Open source code for AlphaFold 2.
PIX is an image processing library in JAX, for JAX.
Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python ⚡
Code for the paper: Distilling the Knowledge from Conditional Normalizing Flows
Probabilistic reasoning and statistical analysis in TensorFlow
Reverb is an efficient and easy-to-use data storage and transport system designed for machine learning research
A place to submit conda recipes before they become fully fledged conda-forge feedstocks
Machine learning algorithms for many-body quantum systems