趋近智
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 章测验
章节 1: JAX 简介
本章介绍 JAX，它是一个为高性能数值计算设计的 Python 库，尤其适用于机器学习研究。我们会首先说明 JAX 是什么，以及它在 Python 科学计算体系中的位置。
你将了解 JAX 的核心思想，重点说明它与 NumPy 的关系以及对函数转换的依赖。我们将比较 jax.numpy 与标准 NumPy 库，指出你需要注意的主要异同点，例如不可变性。我们还将介绍安装 JAX 的具体步骤，以及如何为 CPU、GPU 和 TPU 等不同硬件进行配置。
最后，你将亲手实践创建和操作 JAX 数组（这一核心数据结构），并学习 JAX 如何在可用硬件设备上管理计算。到本章结束时，你将对 JAX 的作用、其基本 API 有一个初步认识，并知道如何设置环境来开始使用它。
课程章节
1.1 JAX 是什么？
1.2 JAX 对比 NumPy
1.3 核心设计理念：函数变换
1.4 安装与设置
1.5 使用 JAX 数组
1.6 设备管理：CPU、GPU、TPU
1.7 动手练习：基本数组操作
© 2026 ApX Machine Learning用心打造