flax 开源项目推荐（358 个）

probml /pyprobml

Python code for "Probabilistic Machine learning" book by Kevin Murphy

FlaxEngine /FlaxEngine

Flax Engine – multi-platform 3D game engine

phlippe /uvadlc_notebooks

阿姆斯特丹大学（MSc AI）教授深度学习课程的Inspyter笔记本教程库，2023年秋季

young-geng /EasyLM

Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.

yang-song /score_sde

通过随机方程进行基于分数的生成式建模的官方代码（ICLR 2021，Oral）

Toni-SM /skrl

模块化强化学习（RL）库（在PyTorch、JAX和NVIDIA Warp中实现），支持Gymnasium/Gym、NVIDIA Isaac Lab、MuJoCo Playground和其他环境

awslabs /fortuna

不确定性量化库。

gordicaleksa /get-started-with-JAX

The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.

google-research /long-range-arena

Long Range Arena for Benchmarking Efficient Transformers

ikostrikov /jaxrl

JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.

texttron /tevatron

Tevatron - Unified Document Retrieval Toolkit across Scale, Language, and Modality. Demo in SIGIR 2023, SIGIR 2025.

google-research /prompt-tuning

Original Implementation of Prompt Tuning from Lester, et al, 2021

google /orbax

Orbax为JAX用户提供常见的检查点和持久性实用程序

luyug /GradCache

Run Effective Large Batch Contrastive Learning Beyond GPU/TPU Memory Constraint

google /CommonLoopUtils

CLU lets you write beautiful training loops in JAX.

HenryNdubuaku /nanodl

JAX library for training sub-4B foundation models for edge

matthias-wright /flaxmodels

Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.

EmptyJackson /unifloral

Unified Implementations of Offline Reinforcement Learning Algorithms

Sea-Snell /JAXSeq

Train very large language models in Jax.

nissymori /JAX-CORL

Clean single-file implementation of offline RL algorithms in JAX