令人惊叹的 JAX

JAX 将自动微分和 XLA 编译器 通过一个类似 NumPy 的 API 结合在一起，用于在 GPU 和 TPU 等加速器上进行高性能机器学习研究。

这是一个精心整理的关于 JAX 的优秀库、项目及其他资源的列表。欢迎贡献！

目录

• 库
• 模型和项目
• 视频
• 论文https://github.com/jax-ml/jax
• 教程和博客文章
• 书籍
• 社区

库

• 神经网络库
  • Flax - 以灵活性和清晰性为核心。
  • Flax NNX - 同一团队对 Flax 的进一步发展
  • Haiku - 专注于简洁性，由 DeepMind 的 Sonnet 团队开发。
  • Objax - 具有类似于 PyTorch 的面向对象设计。
  • Elegy - 用于 JAX 中深度学习的高级 API。支持 Flax、Haiku 和 Optax。
  • Trax - “开箱即用”的深度学习库，专注于提供常见工作负载的解决方案。
  • Jraph - 轻量级图神经网络库。
  • Neural Tangents - 用于指定有限及_无限_宽度神经网络的高级 API。
  • HuggingFace Transformers - 面向广泛自然语言处理任务的预训练 Transformer 生态系统（基于 Flax）。
  • Equinox - 可调用的 PyTree 结构与过滤后的 JIT/梯度变换 => 在 JAX 中构建神经网络。
  • Scenic - 用于计算机视觉研究及其他领域的 JAX 库。
  • Penzai - 优先考虑可读性、可视化以及使用组合工具和简单思维模型轻松编辑神经网络模型。
• Levanter - 使用命名张量和 JAX 构建可读、可扩展、可复现的基础模型。
• EasyLM - 让大型语言模型变得简单：在 JAX/Flax 中进行预训练、微调、评估和部署。
• NumPyro - 基于 Pyro 库的概率编程。
• Chex - 用于编写和测试可靠 JAX 代码的实用工具。
• Optax - 梯度处理与优化库。
• RLax - 用于实现强化学习智能体的库。
• JAX, M.D. - 加速且可微分的分子动力学模拟。
• Coax - 轻松将强化学习论文转化为代码。
• Distrax - TensorFlow Probability 的重新实现，包含概率分布和双射变换。
• cvxpylayers - 构建可微分的凸优化层。
• TensorLy - 让张量学习变得简单。
• NetKet - 用于量子物理的机器学习工具箱。
• Fortuna - AWS 提供的用于深度学习中不确定性量化的一套库。
• BlackJAX - 适用于 JAX 的采样器库。
• Dynamax - 概率状态空间模型。

新兴库

本节收录了一些制作精良且实用，但尚未经过大规模用户群体考验的库。

• 神经网络库
  • FedJAX - 基于Optax和Haiku构建的JAX联邦学习框架。
  • Equivariant MLP - 用于构建等变神经网络层。
  • jax-resnet - Flax中ResNet变体的实现与检查点。
  • jax-raft - RAFT光流估计器的JAX/Flax移植版。
  • Parallax - 适用于JAX的不可变Torch模块。
• 非线性优化
  • Optimistix - 根求解、最小化、不动点及最小二乘法。
  • JAXopt - JAX中的硬件加速（GPU/TPU）、可批处理且可微分的优化器。
• jax-unirep - 实现用于蛋白质机器学习应用的UniRep模型的库。
• flowjax - 以equinox模块形式构建的概率分布与归一化流。
• flaxdiff - 用于在多节点、多设备分布式环境中（如TPU）构建和训练扩散模型的框架与库。
• jax-flows - JAX中的归一化流。
• sklearn-jax-kernels - 使用JAX计算scikit-learn核矩阵。
• jax-cosmo - 可微宇宙学库。
• efax - JAX中的指数族。
• mpi4jax - 将MPI操作与CPU和GPU上的Jax代码结合。
• imax - 图像增强与变换工具。
• FlaxVision - Flax版本的TorchVision。
• Oryx - 基于程序变换的概率编程语言。
• Optimal Transport Tools - 用于解决最优传输问题的工具箱。
• delta PV - 具有自动微分功能的光伏模拟器。
• jaxlie - 用于刚体变换与优化的李理论库。
• BRAX - 可微物理引擎，用于模拟环境以及训练针对这些环境的智能体算法。
• flaxmodels - Jax/Flax的预训练模型。
• CR.Sparse - XLA加速的稀疏表示与压缩感知算法。
• exojax - 兼容JAX的系外行星/棕矮星自动可微谱建模工具。
• PIX - JAX中的图像处理库。
• bayex - 基于JAX的贝叶斯优化。
• JaxDF - 用于任意离散化的可微仿真框架。
• tree-math - 将作用于数组的函数转换为作用于PyTree的函数。
• jax-models - 原本无代码或使用非JAX框架编写的科研论文实现。
• PGMax - 用于构建离散概率图模型（PGM）并在JAX上进行推理的框架。
• EvoJAX - 硬件加速的神经进化。
• evosax - 基于JAX的进化策略。
• SymJAX - 符号化的CPU/GPU/TPU编程。
• mcx - 表达并编译概率程序以实现高效推理。
• Einshape - 基于DSL的重塑库，适用于JAX及其他框架。
• ALX - 开源分布式矩阵分解库，采用交替最小二乘法，更多信息请参见ALX: Large Scale Matrix Factorization on TPUs。
• Diffrax - JAX中的数值微分方程求解器。
• tinygp - JAX中最“小巧”的高斯过程库。
• gymnax - 具有知名gym API的强化学习环境。
• Mctx - 原生JAX中的蒙特卡洛树搜索算法。
• KFAC-JAX - 用于神经网络的近似曲率二阶优化。
• TF2JAX - 将函数/图转换为JAX函数。
• jwave - 用于可微声学模拟的库。
• GPJax - JAX中的高斯过程。
• Jumanji - 一套由行业驱动、基于JAX编写的硬件加速RL环境。
• Eqxvision - Equinox版本的TorchVision。
• JAXFit - 加速曲线拟合库，用于非线性最小二乘问题（参见arXiv论文）。
• econpizza - 使用JAX求解具有异质性主体的宏观经济模型。
• SPU - 用于运行带有MPC（安全多方计算）的JAX代码的领域特定编译器和运行时套件。
• jax-tqdm - 为JAX扫描和循环添加tqdm进度条。
• safejax - 使用🤗safetensors序列化JAX、Flax、Haiku或Objax模型参数。
• Kernex - JAX中的可微模板装饰器。
• MaxText - 用纯Python/Jax编写、面向Google Cloud TPU的简单、高效且可扩展的Jax LLM。
• Pax - 基于JAX的大规模模型训练机器学习框架。
• Praxis - Pax的层库，旨在供其他基于JAX的机器学习项目使用。
• purejaxrl - JAX中的可向量化、端到端强化学习算法。
• Lorax - 自动将LoRA应用于JAX模型（Flax、Haiku等）。
• SCICO - JAX中的科学计算成像。
• Spyx - JAX中的脉冲神经网络，用于类脑硬件上的机器学习。
• 脑动力学编程生态系统
  • BrainPy - Python中的脑动力学编程。
  • brainunit - JAX中的物理单位与单位感知数学系统。
  • dendritex - JAX中的树突建模。
  • brainstate - 基于状态的程序编译与增强系统。
  • braintaichi - 利用Taichi语言自定义脑动力学算子。
• OTT-JAX - JAX中的最优传输工具。
• QDax - Jax中的质量多样性优化。
• JAX Toolbox - 每日CI及针对NVIDIA GPU上JAX的优化示例，使用T5x、Paxml和Transformer Engine等库。
• Pgx - 向量化的棋盘游戏环境，用于强化学习，包含AlphaZero示例。
• EasyDeL - EasyDeL 🔮 是一个开源库，旨在通过酷炫的训练与服务选项（如Llama、MPT、Mixtral、Falcon等）使您的训练更快、更优化。
• XLB - 用Python编写的可微大规模并行格子玻尔兹曼库，用于基于物理的机器学习。
• dynamiqs - 使用JAX进行高性能且可微的量子系统模拟。
• foragax - 基于JAX的代理建模框架。
• tmmax - 使用JAX对薄膜结构光学性质进行向量化计算。是薄膜光学研究的瑞士军刀式工具。
• Coreax - 用于寻找核集以压缩大型数据集同时保留其统计特性的小型算法库。
• NAVIX - 在JAX中重新实现的MiniGrid强化学习环境。
• FDTDX - JAX中的有限差分时域电磁仿真。
• DiffeRT - 基于JAX生态系统的可微射线追踪工具箱，用于无线电传播。
• JAX-in-Cell - 等离子体物理模拟，采用PIC（粒子在细胞）方法自洽地求解电磁场中的电子和离子动力学。
• kvax - JAX中的FlashAttention实现，支持高效的文档掩码计算和上下文并行性。
• astronomix - JAX中的天体物理学可微（磁流体）动力学。
• vivsim - 使用浸入边界-格子玻尔兹曼方法进行流固耦合模拟。
• MBIRJAX - 高性能断层扫描重建。
• torchax - torchax是一个用于Jax与PyTorch编写模型代码互操作的库。

模型与项目

JAX

• Fourier Feature Networks - 官方实现，对应论文 Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains。
• kalman-jax - 使用迭代卡尔曼滤波和平滑算法对马尔可夫（即时间序列）高斯过程进行近似推断。
• jaxns - 在 JAX 中实现嵌套采样。
• Amortized Bayesian Optimization - 与论文 离散空间上的摊销贝叶斯优化 相关的代码。
• Accurate Quantized Training - 用于在 JAX 和 Flax 中运行及分析神经网络量化实验的工具和库。
• BNN-HMC - 论文 贝叶斯神经网络后验分布究竟是什么样的？ 的实现。
• JAX-DFT - 在 JAX 中实现的一维密度泛函理论（DFT），包括论文 Kohn-Sham 方程作为正则化项：将先验知识融入机器学习物理模型 的内容。
• Robust Loss - 论文 一种通用且自适应的鲁棒损失函数 的参考代码。
• Symbolic Functionals - 来自论文 演化符号化的密度泛函 的演示。
• TriMap - 论文 TriMap: 基于三元组的大规模降维 的官方 JAX 实现。

Flax

• awesome-jax-flax-llms - 用 JAX 和 Flax 实现的大规模语言模型合集
• DeepSeek-R1-Flax-1.5B-Distill - DeepSeek-R1 1.5B 精炼推理大模型的 Flax 实现。
• Performer - Performer（通过 FAVOR+ 实现的线性 Transformer）架构的 Flax 实现。
• JaxNeRF - 基于多设备 GPU/TPU 支持，实现 NeRF：以神经辐射场表示场景用于视图合成。
• mip-NeRF - 官方实现 Mip-NeRF：一种用于抗锯齿的多尺度神经辐射场表示。
• RegNeRF - 官方实现 RegNeRF：针对稀疏输入的视图合成中正则化神经辐射场。
• JaxNeuS - 实现 NeuS：通过体积渲染学习神经隐式表面用于多视角重建。
• Big Transfer (BiT) - 实现 Big Transfer (BiT)：通用视觉表征学习。
• JAX RL - 强化学习算法的实现。
• gMLP - 实现 关注 MLP。
• MLP Mixer - 极简实现 MLP-Mixer：一种全 MLP 的视觉架构。
• Distributed Shampoo - 实现 二阶优化的实用化。
• NesT - 官方实现 聚合嵌套 Transformer。
• XMC-GAN - 官方实现 跨模态对比学习用于文本到图像生成。
• FNet - 官方实现 FNet：用傅里叶变换混合 Token。
• GFSA - 官方实现 利用有限状态自动机层学习图结构。
• IPA-GNN - 官方实现 通过指令指针注意力图神经网络学习执行程序。
• Flax Models - Flax 中实现的各种模型和方法合集。
• Protein LM - 实现 BERT 和自回归模型用于蛋白质，如 生物结构与功能源于将无监督学习扩展至 2.5 亿条蛋白质序列 和 ProGen：用于蛋白质生成的语言模型 所述。
• Slot Attention - 参考实现 用于图像识别的可微分补丁选择。
• Vision Transformer - 官方实现 一张图胜过 16×16 个词：大规模图像识别中的 Transformer。
• FID 计算 - 将 mseitzer/pytorch-fid 移植到 Flax。
• ARDM - 官方实现 自回归扩散模型。
• D3PM - 官方实现 离散状态空间中的结构化去噪扩散模型。
• Gumbel-max 因果机制 - 用于 学习广义 Gumbel-max 因果机制 的代码，额外代码见 GuyLor/gumbel_max_causal_gadgets_part2。
• Latent Programmer - 用于 ICML 2021 论文 Latent Programmer：用于程序合成的离散潜在编码 的代码。
• SNeRG - 官方实现 烘焙神经辐射场用于实时视图合成。
• 自旋加权球面 CNN - 对 自旋加权球面 CNN 的改编。
• VDVAE - 对 非常深的 VAE 能够泛化自回归模型，并在图像上超越它们 的改编，原始代码见 openai/vdvae。
• MUSIQ - ICCV 2021 论文 MUSIQ：多尺度图像质量 Transformer 的检查点及模型推理代码。
• AQuaDem - 官方实现 基于演示的动作量化连续控制。
• Combiner - 官方实现 Combiner：全注意力 Transformer，计算成本稀疏。
• Dreamfields - ICLR 2022 论文 用于快速采样扩散模型的渐进蒸馏 的官方实现。
• GIFT - 官方实现 野外渐进式领域适应：当中间分布不存在时。
• 光场神经渲染 - 官方实现 光场神经渲染。
• Raphael Pisoni 在 JAX 中实现的锐化余弦相似度 - 锐化余弦相似度层的 JAX/Flax 实现。
• 用于求解组合优化问题的图神经网络 - 基于 JAX + Flax 的 受物理启发的图神经网络进行组合优化 的实现。
• DETR - 使用 Sinkhorn 求解器和并行二部匹配，实现 DETR：端到端目标检测 Transformer 的 Flax 版本。

Haiku

• AlphaFold - AlphaFold v2.0 推理流程的实现，相关论文为 Highly accurate protein structure prediction with AlphaFold。
• Adversarial Robustness - 参考代码，用于 Uncovering the Limits of Adversarial Training against Norm-Bounded Adversarial Examples 和 Fixing Data Augmentation to Improve Adversarial Robustness 两篇论文。
• Bootstrap Your Own Latent - 论文 Bootstrap your own latent: A new approach to self-supervised Learning 的实现。
• Gated Linear Networks - GLN 是一类无需反向传播的神经网络。
• Glassy Dynamics - 论文 Unveiling the predictive power of static structure in glassy systems 的开源实现。
• MMV - 用于 Self-Supervised MultiModal Versatile Networks 中模型的代码。
• Normalizer-Free Networks - 论文 NFNets 的官方 Haiku 实现。
• NuX - 基于 JAX 的归一化流。
• OGB-LSC - 该仓库包含 DeepMind 参与 OGB Large-Scale Challenge 的 PCQM4M-LSC（量子化学）和 MAG240M-LSC（学术图）赛道的提交。
• Persistent Evolution Strategies - 用于论文 Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies 的代码。
• Two Player Auction Learning - 论文 Auction learning as a two-player game 的 JAX 实现。
• WikiGraphs - 用于复现 WikiGraphs: A Wikipedia Text - Knowledge Graph Paired Dataset 研究结果的基准代码。

Trax

• Reformer - Reformer（高效 Transformer）架构的实现。

NumPyro

• lqg - 论文 Putting perception into action with inverse optimal control for continuous psychophysics 中针对线性二次高斯问题的贝叶斯逆最优控制的官方实现。

Equinox

• Sampling Path Candidates with Machine Learning - 论文 Towards Generative Ray Path Sampling for Faster Point-to-Point Ray Tracing 的官方教程及实现。

视频

• NeurIPS 2020: JAX 生态系统聚会 - 讨论 JAX、其在 DeepMind 的应用，以及工程师、科学家与 JAX 核心团队之间的交流。
• JAX 入门 - 使用 JAX 从零开始构建简单神经网络。
• JAX：加速机器学习研究 | SciPy 2020 | VanderPlas - JAX 的核心设计、如何推动新研究，以及如何开始使用它。
• 使用 JAX + NumPyro 进行贝叶斯编程 — Andy Kitchen - 使用 NumPyro 进行贝叶斯建模的入门介绍。
• JAX：通过 Python 中的可组合函数变换加速机器学习研究 | NeurIPS 2019 | Skye Wanderman-Milne - 在 Program Transformations for Machine Learning 工作坊中的 JAX 入门演讲。
• JAX 在云端 TPU 上 | NeurIPS 2020 | Skye Wanderman-Milne 和 James Bradbury - 演示 TPU 主机访问的报告。
• 深度隐层——神经 ODE、深度均衡模型及其他 | NeurIPS 2020 - 由 Zico Kolter、David Duvenaud 和 Matt Johnson 制作的教程，附带 Colab 笔记本，可在 Deep Implicit Layers 中找到。
• 使用 JAX 在 TPU Pod 片段上求解 y=mx+b - Mat Kelcey - 一个包含四集 YouTube 教程的系列，配有 Colab 笔记本，从 JAX 基础知识开始，逐步过渡到在 v3-32 TPU Pod 片段上采用数据并行方式进行训练。
• JAX、Flax 和 Transformers 🤗 - 围绕 JAX / Flax、Transformers、大规模语言建模等主题的为期三天的讲座。

论文

本节包含专注于 JAX 的论文（例如基于 JAX 的库白皮书、关于 JAX 的研究等）。使用 JAX 实现的论文列在 模型/项目 部分中。

• 通过高级跟踪编译机器学习程序. Roy Frostig、Matthew James Johnson、Chris Leary。MLSys 2018。 - 白皮书描述了 JAX 的早期版本，详细介绍了计算如何被跟踪和编译。
• JAX, M.D.: 可微分物理框架. Samuel S. Schoenholz、Ekin D. Cubuk。NeurIPS 2020。 - 介绍了 JAX, M.D.，一个包含仿真环境、相互作用势、神经网络等内容的可微分物理库。
• 通过即时编译与向量化实现快速差分隐私 SGD. Pranav Subramani、Nicholas Vadivelu、Gautam Kamath。arXiv 2020。 - 利用 JAX 的 JIT 和 VMAP，实现了比现有库更快的差分隐私计算。
• XLB：Python 中的可微分大规模并行格子玻尔兹曼库. Mohammadmehdi Ataei、Hesam Salehipour。arXiv 2023。 - 白皮书描述了 XLB 库：基准测试、验证以及关于该库的更多细节。

教程和博客文章

• 使用 JAX 加速研究，作者：David Budden 和 Matteo Hessel - 描述了 DeepMind 中 JAX 及其生态系统的现状。
• JAX 入门（MLP、CNN 和 RNN），作者：Robert Lange - 使用 JAX 的基础运算符从头构建神经网络模块。
• 学习 JAX：从线性回归到神经网络，作者：Rito Ghosh - 对 JAX 进行温和的介绍，并利用它实现线性和逻辑回归以及神经网络模型，进而解决实际问题。
• 教程：使用 JAX 和 Flax Linen 进行图像分类，作者：8bitmp3 - 学习如何使用 Flax 的 Linen API 创建一个简单的卷积网络，并训练它来识别手写数字。
• 深入 JAX，作者：Nick Doiron - 在 Kaggle 花卉分类挑战中比较了 Flax、Haiku 和 Objax。
• 50 行代码实现元学习，作者：Eric Jang - 同时介绍了 JAX 和元学习。
• 100 行代码实现归一化流，作者：Eric Jang - 简洁地实现了 RealNVP。
• 在 GPU/TPU 上进行可微路径追踪，作者：Eric Jang - 实现路径追踪的教程。
• 集成网络，作者：Mat Kelcey - 集成网络是一种将多个模型的集合表示为单个逻辑模型的方法。
• 分布外检测，作者：Mat Kelcey - 实现了多种 OOD 检测方法。
• 使用 JAX 理解自动微分，作者：Srihari Radhakrishna - 通过 JAX 理解自动微分的工作原理。
• 从 PyTorch 到 JAX：迈向净化有状态代码的神经网络框架，作者：Sabrina J. Mielke - 展示了如何从类似 PyTorch 的编码风格过渡到更函数式的编码风格。
• 使用自定义 C++ 和 CUDA 代码扩展 JAX，作者：Dan Foreman-Mackey - 教程演示了在 JAX 中提供自定义操作所需的基础设施。
• 在 JAX 中进化神经网络，作者：Robert Tjarko Lange - 探讨了 JAX 如何推动下一代可扩展的神经进化算法的发展。
• 使用 JAX 探索超参数元损失景观，作者：Luke Metz - 演示如何使用 JAX 进行 SGD 和动量法的内层损失优化、基于梯度的外层损失优化，以及使用进化策略进行外层损失优化。
• 在 JAX 中实现确定性 ADVI，作者：Martin Ingram - 详细讲解如何使用 JAX 轻松干净地实现自动微分变分推断 (ADVI)。
• 进化通道选择，作者：Mat Kelcey - 训练一个对不同分辨率下输入通道组合具有鲁棒性的分类模型，然后使用遗传算法来决定针对特定损失的最佳组合。
• JAX 入门，作者：Kevin Murphy - Colab 介绍了该语言的各个方面，并将其应用于简单的机器学习问题。
• 在 JAX 中编写 MCMC 采样器，作者：Jeremie Coullon - 教程介绍了在 JAX 中编写 MCMC 采样器的不同方式，并附带速度基准测试。
• 如何在 JAX 的扫描和循环中添加进度条，作者：Jeremie Coullon - 教程说明如何使用 host_callback 模块在 JAX 的编译循环中添加进度条。
• JAX 入门，作者：Aleksa Gordić - 一系列笔记本和视频，从零基础开始逐步学习 JAX，最终在 Haiku 中构建神经网络。
• 使用 JAX + FLAX 编写训练循环，作者：Saurav Maheshkar 和 Soumik Rakshit - 教程介绍了如何在 JAX、Flax 和 Optax 中编写一个简单、端到端的训练和评估流水线。
• 在 JAX 中实现 NeRF，作者：Soumik Rakshit 和 Saurav Maheshkar - 教程介绍了如何在 JAX 中使用神经辐射场表示场景并进行三维体积渲染。
• 使用 JAX+Flax 的深度学习教程，作者：Phillip Lippe - 一系列笔记本，从基础知识（如 JAX/Flax 入门、激活函数）到最新进展（如 Vision Transformer、SimCLR），都配有 PyTorch 版本。
• 通过 PureJaxRL 实现 4000 倍加速，作者：Chris Lu - 一篇关于如何通过向量化使 JAX 极大地加速强化学习训练的博客文章。
• 使用 JAX 编写简单的 PDE 求解器 + 约束优化，作者：Philip Mocz - 一个简单的例子，展示了如何使用 JAX 解决对流-扩散方程，并将其用于约束优化问题，以找到产生预期结果的初始条件。

书籍

• JAX 实战 - 一本关于如何使用 JAX 进行深度学习和其他数学密集型应用的实践指南。

社区

• JaxLLM（非官方）Discord
• JAX GitHub 讨论区
• Reddit

贡献

欢迎贡献！请先阅读贡献指南。

Awesome JAX 快速上手指南

JAX 是 Google 推出的高性能机器学习研究框架，它将自动微分（Automatic Differentiation）与 XLA 编译器结合，提供了类似 NumPy 的 API，能够高效利用 GPU 和 TPU 加速器。本指南将帮助你快速搭建环境并开始使用 JAX 生态。

1. 环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统: Linux (推荐), macOS, 或 Windows (WSL2 推荐)。
• Python 版本: Python 3.9 - 3.12。
• 硬件加速:
  • GPU: 需要安装 NVIDIA 驱动和 CUDA toolkit (JAX 通常自带兼容的 CUDA/cuDNN 版本，但需确保驱动较新)。
  • TPU: 需要在 Google Colab、Kaggle 或 Google Cloud TPU VM 上使用。
• 前置依赖: 建议先更新 pip 并安装基础构建工具。

python -m pip install --upgrade pip

国内开发者提示：如果遇到网络连接问题导致下载缓慢，建议使用国内镜像源（如清华源、阿里源）进行安装。

2. 安装步骤

根据你的硬件环境选择以下一种安装方式。

方案 A：CPU 版本（通用）

适用于没有独立显卡或仅用于调试代码的场景。

pip install jax jaxlib

方案 B：GPU 版本（NVIDIA）

适用于本地拥有 NVIDIA 显卡的开发环境。此命令会自动安装与当前 JAX 版本匹配的 CUDA 和 cuDNN 依赖。

pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

国内加速安装命令（推荐）： 使用清华镜像源加速 Python 包下载，同时指定 JAX CUDA 源：

pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html -i https://pypi.tuna.tsinghua.edu.cn/simple

方案 C：TPU 版本

适用于 Google Colab 或 TPU VM 环境。

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

安装常用生态库

Awesome JAX 列表中包含了丰富的生态系统库，最常用的组合是 Flax (神经网络库) 和 Optax (优化器库)：

pip install flax optax orbax-checkpoint -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 基本使用

JAX 的核心设计理念是函数式编程。以下是最简单的示例，展示如何使用 jit 加速计算以及如何使用 grad 进行自动求导。

示例：加速计算与自动微分

创建一个名为 demo.py 的文件：

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap

# 1. 定义一个普通的 Python 函数 (使用 jax.numpy 替代 numpy)
def sigmoid(x):
    return 1 / (1 + jnp.exp(-x))

# 2. 使用 jit 编译加速 (Just-In-Time Compilation)
# 第一次运行会编译，后续运行速度极快
fast_sigmoid = jit(sigmoid)

# 3. 使用 grad 自动计算梯度
grad_sigmoid = grad(sigmoid)

# 4. 执行计算
x = jnp.array([0.0, 1.0, -1.0])
y = fast_sigmoid(x)
dy_dx = grad_sigmoid(x)

print(f"输入：{x}")
print(f"Sigmoid 输出：{y}")
print(f"梯度输出：{dy_dx}")

# 5. (可选) 使用 vmap 自动向量化
# 假设我们要对一批形状为 (batch, features) 的数据操作
batch_x = jnp.ones((5, 3))
# vmap 会自动在第一个维度上广播 sigmoid 函数
batch_y = vmap(sigmoid)(batch_x) 
print(f"批量计算输出形状：{batch_y.shape}")

运行脚本：

python demo.py

下一步探索

根据 Awesome JAX 列表，你可以进一步尝试：

• 构建模型: 使用 Flax 或 Haiku 构建神经网络。
• 优化训练: 使用 Optax 配置优化器（如 AdamW）。
• 概率编程: 尝试 NumPyro 进行贝叶斯推断。
• 强化学习: 使用 RLax 或 Brax 进行物理仿真与训练。