趋近智
JAX 入门
章节 1: JAX 简介
JAX 是什么？
JAX 对比 NumPy
核心设计理念：函数变换
安装与设置
使用 JAX 数组
设备管理：CPU、GPU、TPU
动手练习：基本数组操作
第 1 章测验
章节 2: 通过 JIT 编译加速函数
速度提升：为何需要编译？
介绍
jax.jitJIT 工作原理：追踪与编译
Python 控制流与
jit静态值与跟踪值
jit 的常见问题动手实践：应用
jit第 2 章测验
章节 3: 使用
grad 进行自动微分理解梯度
介绍
jax.grad自动微分的工作方式：反向模式
关于参数求导
高阶导数（
grad的grad）值和梯度 (
jax.value_and_grad)求导与控制流
局限性与注意事项
动手实践：计算梯度
第 3 章测验
章节 4: 使用
vmap 实现自动向量化向量化的原理
介绍
jax.vmap对特定参数进行映射（
in_axes，out_axes）处理多个批处理参数
嵌套
vmap结合
vmap 与 jit 和 gradvmap的性能考量动手实践：函数向量化
第 4 章测验
章节 5: 使用
pmap 在多设备上并行计算数据并行 (SPMD) 介绍
介绍
jax.pmap将数据映射到设备 (
in_axes, out_axes)设备网格与轴名称
集体操作（
lax.psum、lax.pmean等）将
pmap 与其他变换结合使用调试
pmap 化的函数动手实践：并行计算
第 5 章测验
章节 6: JAX 中的状态管理
函数纯粹性与副作用
函数式代码中的状态挑战
模式：显式状态传递
使用 PyTree 管理分层状态
示例：有状态计数器
例子：简单的优化器状态
将状态管理与变换结合
实践：实现有状态函数
第 6 章测验
章节 2: 通过 JIT 编译加速函数
尽管 JAX 提供了一个熟悉的 NumPy 风格的 API，但对于科学计算和机器学习中常见的密集数值计算，标准的 Python 执行速度可能较慢，尤其是在针对 GPU 或 TPU 等加速器时。为了获得高效率，JAX 依赖于编译。
本章主要讲解 jax.jit，这是一个即时 (JIT) 编译转换。您将了解 jit 如何将 Python 函数转换成针对您的硬件进行适配的优化后的可执行代码。我们将介绍：
jax.jit作为装饰器或函数的基本使用。- 即时 (JIT) 编译如何通过一个称为“追踪”的过程来工作。
- 追踪对 Python 控制流（循环、条件语句）的影响。
- 在编译函数中如何处理静态值与动态值。
- 使用
jit时的常见错误和注意事项。
在本章结束时，您将能够使用 jax.jit 大幅提升您的 JAX 计算速度。
课程章节
2.1 速度提升：为何需要编译？
2.2 介绍 `jax.jit`
2.3 JIT 工作原理：追踪与编译
2.4 Python 控制流与 `jit`
2.5 静态值与跟踪值
2.6 `jit` 的常见问题
2.7 动手实践：应用 `jit`
© 2026 ApX Machine Learning用心打造