内容简介

《JAX可微分编程》以Google开发的JAX开源框架为载体，详细介绍了JAX在可微分编程领域的应用，具体包括自动微分的基本原理、数据结构，以及自动微分在实际场景中的应用，其涉及的领域包括但不限于算法优化、神经网络、工程建模、量子计算等。

《JAX可微分编程》分为3部分，总计10章外加5篇附录。第1部分介绍了可微分编程的基本原理，包括手动求导、数值微分、符号微分以及自动微分的前向模式和反向模式，在未调用任何库函数的情况下，从零开始构建起了符号微分及自动微分的数据结构。第2部分是对JAX库特性的介绍，包括JAX的基本语法、自动微分、即时编译和并行计算，并以此为切口，对深度学习、λ演算等领域进行了深入浅出的讲解。第3部分是自动微分在实际场景中的应用，除了算法优化、神经网络等极其经典的应用场景，这一部分还给出了自动微分在工程建模、量子计算等方面的理论及应用。

《JAX可微分编程》涵盖的范围虽广，但对其中每个领域的介绍都绝非浅尝辄止，无论材料的选取、内容的编排，还是论述的视角、观点的呈现，均不乏新颖之处。通过本书的学习，读者不但可以掌握JAX开源框架的用法，还可以获悉JAX在可微分编程领域的具体应用方法。本书适合在工作中会用到自动微分技术的工程技术人员、高校科研人员阅读，也适合对JAX感兴趣并希望能掌握其应用的AI技术人员阅读。

作品目录

六重奏女士的诅咒 1
第1章　程序视角下的微分运算 9
1.1　函数与求导 9
1.1.1 求导的基本概念 10
1.1.2 梯度操作（Gradient Operator） 11
1.1.3 雅可比矩阵（Jacobian Matrix） 12
1.1.4 黑塞矩阵（Hessian Matrix） 13
1.2　手动求导 14
1.3　数值微分（Numeric Differentiation） 14
1.3.1 数值微分的理论基础 14
1.3.2 数值微分的误差来源 16
1.3.3 数值微分的程序实现 19
1.4　符号微分（Symbolic Differentiation） 24
1.4.1 计算图 24
1.4.2 计算图的构建 26
1.4.3 SymPy库简介 34
第2章　自动微分 39
2.1　前向模式（forward mode） 40
2.1.1 前向模式的理论 40
*2.1.2 前向模式的二元数诠释 44
2.1.3 前向微分的程序实现 49
2.2　反向模式（backward mode） 58
2.2.1 反向模式的理论 58
*2.2.2 反向模式和前向模式的统一 60
2.2.3 反向模式的程序实现 67
第3章　初识JAX 78
3.1　数组的创建 79
3.1.1 数组的性质 79
3.1.2 创建数组的函数 81
3.1.3 随机数组的创建 84
3.2　数组的修改 86
3.2.1 多维数组的重排 86
3.2.2 多维数组的扩展 89
3.2.3 多维数组的索引 92
3.2.4 越界行为的处理 95
3.2.5 异地更新 95
3.3　数组的运算 97
3.3.1 语义广播 97
3.3.2 数组运算 98
3.3.3 线性代数 100
3.3.4 科学计算 102
3.3.5 爱因斯坦求和约定 102
3.4　使用蒙特卡罗方法估计圆周率 105
第4章　JAX的微分运算 109
4.1　微分操作的语法 109
4.1.1 JAX中的梯度操作 110
4.1.2 JAX中的雅可比矩阵 116
4.1.3 JAX中的黑塞矩阵 118
4.1.4 自定义算符及隐函数求导 120
4.2　梯度下降 125
4.2.1 从最小二乘法说开去 126
4.2.2 寻找极小值 130
4.2.3 训练及误差 133
4.2.4 全连接神经网络 140
第5章　JAX的编程范式及即时编译 150
5.1　函数式编程 151
5.1.1 函数式编程的价值观 151
5.1.2 JAX中的即时编译 158
5.1.3 JAX中的条件语句 162
5.1.4 JAX中的流程控制语句 166
5.1.5 静态变量 170
*5.2　λ演算 172
5.2.1 λ演算的基本设定 172
5.2.2 λ演算中的布尔代数 176
5.2.3 λ演算中的自然数 177
5.2.4 λ演算中的递归 181
第6章　JAX的并行计算 186
6.1　函数vmap 187
6.2　使用vmap在GPU上并行训练 194
6.3　函数pmap 197
6.4　使用pmap更新细胞自动机 200
第7章　优化算法 208
7.1　下降算法概要 209
7.1.1 下降算法的数学表述 209
7.1.2 步长的选取 213
7.1.3 终止条件的选取 218
7.1.4 下降方向的选取 219
*7.1.5 共轭梯度算法 223
7.2　一阶优化算法 233
7.2.1 动量法 233
7.2.2 自适应算法 242
7.2.3 Adam 246
第8章　循环神经网络 250
*8.1　神经网络的生物学基础 250
8.1.1 神经元的电化学性质 251
8.1.2 神经元输出过程的建模 254
8.1.3 神经元构成网络的建模 256
8.2　循环神经网络 263
8.2.1 简单循环神经网络 263
8.2.2 循环神经网络的梯度回传 272
8.2.3 简单循环神经网络的程序实现 277
8.2.4 长短期记忆单元及其程序实现 286
8.2.5 案例：股票预测 292
第9章　案例：FAST主动反射面的形态调节 298
9.1　背景介绍 298
9.2　数据的预处理 301
9.3　约束优化问题的提出及模型的训练 306
9.4 程序运行结果的讨论 314
第10章　量子计算中的自动微分 316
*10.1 量子计算的数学基础 317
10.1.1 算符与量子态 317
10.1.2 算符的指数 323
10.1.3 算符的对易子 326
*10.2 量子计算的物理基础 329
10.2.1 波粒二象性 329
10.2.2 薛定谔方程 331
10.2.3 动量空间 338
*10.3 基于量子体系的自动微分 341
10.3.1 量子比特 342
10.3.2 参数优化 347
附录A　Python中类的介绍 354
附录B　拓扑排序 369
附录C　信息和熵 376
附录D　下降算法的收敛性分析 390
附录E　神经元的Hodgkin Huxley模型 400
后记 410
· · · · · ·

作者简介

程琪皓，北京大学物理学院本科生，曾获第36届全国中学生物理竞赛（浙江赛区）一等奖、“未名学子”奖学金、沈克琦奖学金、北京大学三等奖学金等多项奖励。研究方向涉及强化学习、低维超导实验、量子计算、分子模拟等多个领域。