JAX

开发者发布开源项目，用约 190 行纯 JAX 代码从零实现 DeepSeek V3 核心架构，包含多头潜在注意力（MLA）、解耦 RoPE、DeepSeekMoE 与无辅助损失负载均衡等关键设计，并在 MacBook Air 上完成预训练。项目不依赖 Flax、Equinox 或 PyTorch，参数以纯字典数组表示，代码可逐行阅读，是理解前沿 MoE 架构的高价值教学与实验范本。