jax-llm-examples

纯JAX极简代码，跑通LLM训练与推理

jax-llm-examples 是一个用纯 JAX 实现的极简大语言模型示例集合，目标是让开发者以最少依赖理解并跑通 LLM 的训练与推理全流程。它不依赖 PyTorch 或 TensorFlow，只基于 JAX 及其生态（如 Flax/Optax 风格的手写模块），提供数据加载、模型定义、损失计算、优化器更新和文本生成等完整环节的清晰代码。项目强调性能与可读性的平衡：既展示 JAX 的 jit、vmap、pmap 等加速特性，又保持代码量小、结构直观，适合作为学习 JAX 做 LLM 的入门模板，也便于研究者快速改造实验。相比动辄数千行的训练框架，它把核心逻辑压缩到可通读的规模，帮助用户避开抽象层，直接观察张量流与梯度更新，是理解 LLM 底层实现和 JAX 并行能力的实用参考。