趋近智
JAX 入门
章节 1: JAX 简介
JAX 是什么？
JAX 对比 NumPy
核心设计理念：函数变换
安装与设置
使用 JAX 数组
设备管理：CPU、GPU、TPU
动手练习：基本数组操作
章节 2: 通过 JIT 编译加速函数
速度提升：为何需要编译？
介绍
jax.jitJIT 工作原理：追踪与编译
Python 控制流与
jit静态值与跟踪值
jit 的常见问题动手实践：应用
jit章节 3: 使用
grad 进行自动微分理解梯度
介绍
jax.grad自动微分的工作方式：反向模式
关于参数求导
高阶导数（
grad的grad）值和梯度 (
jax.value_and_grad)求导与控制流
局限性与注意事项
动手实践：计算梯度
章节 4: 使用
vmap 实现自动向量化向量化的原理
介绍
jax.vmap对特定参数进行映射（
in_axes，out_axes）处理多个批处理参数
嵌套
vmap结合
vmap 与 jit 和 gradvmap的性能考量动手实践：函数向量化
章节 5: 使用
pmap 在多设备上并行计算数据并行 (SPMD) 介绍
介绍
jax.pmap将数据映射到设备 (
in_axes, out_axes)设备网格与轴名称
集体操作（
lax.psum、lax.pmean等）将
pmap 与其他变换结合使用调试
pmap 化的函数动手实践：并行计算
章节 6: JAX 中的状态管理
函数纯粹性与副作用
函数式代码中的状态挑战
模式：显式状态传递
使用 PyTree 管理分层状态
示例：有状态计数器
例子：简单的优化器状态
将状态管理与变换结合
实践：实现有状态函数
处理多个批处理参数
这部分内容有帮助吗？
- jax.vmap, JAX developers, 2024 (JAX Documentation) - 提供关于
vmap的最新、准确信息，包括in_axes的用法和示例。 - JAX: Composable transformations of NumPy programs, Roy Frostig, Matthew Johnson, and Christopher Leary, 2018 (Google AI Blog) - 介绍 JAX 的设计，说明
vmap如何融入其数值计算的可组合函数变换系统。 - Autograd: Effortless Automatic Differentiation in Python, Dougal Maclaurin, Michael Johnson, and David Duvenaud, 2015 Proceedings of the International Conference on Machine Learning (ICML), Vol. 37 (PMLR (Proceedings of Machine Learning Research)) DOI: 10.5555/3045118.3045239 - 描述 Autograd，JAX 的前身，阐述了影响 JAX 批处理能力的自动微分和数组编程原理。