jax可微分编程的笔记（3）

第3章初识JAX

JAX是Google开发的高性能数值计算和自动微分库，提供自动微分
即时编译和矢量并行化三大功能。JAX选择将函数编程的思想贯穿
始终。
简单来说，JAX库是GPU加速，支持自动微分的NumPy,NumPy则是
Python之下的基础数值运算库。

3.1 数组的创建

数组的创建是程序书写的开始，在NumPy中，数组的类型是ndarray，
而在JAX中，数组则名为DeviceArray，二者除了名称不同，其使用
方式具有极大的相似之处。

3.1.1 数组的性质

我们可以把标量视为零维数组，将向量视为一维数组，将矩阵视为
二维数组。
数组的性质包括且不仅限于如下的内容：
类型，维度，形状，元素个数，元素的数据类型。

3.1.2 数组的创建函数

zeros(shape[,dtype]) 生成填满0的多维数组

3.1.3 随机数组的创建

exponential(key[,shape,dtype]) 从指数分布中采样指定
形状的随机数组。

3.2 数组的修改

在完成数组的创建后，绝大多数数组的形状的修改都将由数组的重
排、扩展、及索引来完成.

3.2.1 多维数组的重排

reshape(arr,newshape[,order]) 在不改变数组的数据情况下，
改变数组的形状。

3.2.2 多维数组的扩展

repeat(a,repeats[,axis,total_length]) 以重复每个元素
的方式扩展数组。

3.2.3 多维数组的索引

我们可以用方括号[]在数组中索引一个元素，可以用正整数或者是
负整数。还可以用切片操作获得数组的一个部分，语法为[start:end:step]
我们还可以用数组作为索引，来索引多维数组。

3.2.4 越界的行为的处理

在索引元素之前，需要对代码进行必要的检查。

3.2.5 异地更新

原地更新是指直接修改原数组。
异地更新是指修改原数组的副本。

3.3 数组的运算

语义广播在数组的运算中处于基础的地位。

3.3.1 语义的广播

广播描述了在运算的过程中，程序对不同的形状的数组的处理规则。
语义广播的使用条件，当两个数组进行运算时，程序由右向左比较
它们的形状，如果在一个维度上，两个数组的长度相等或者是其中一个
的长度为1，我们认为这两个数组在这个维度上是兼容的。如果两个数组
（从右向左）在每一个维度上都是兼容的，则认为它们之间可以进行
语义广播。在两个数组的维度不同时，将维度小的数组进行向左扩展。

3.3.2 数组的运算

数组的运算的最基础的运算是逐元素进行运算。逐元素运算是一元函
数接受数组作为参数，返回一个新数组。等价于将原本的一元函数作用
于数组中的每个元素。
同理，二元标量运算符同样可以应用于两个数组上，得到一个新的数组。

3.3.3 线性代数的运算

线性代数的运算包括且不仅限于如下的函数：
det(a) 计算方阵a的行列式
eig(a) 计算方阵a的特征值和（右）特征向量
inv(a) 计算方阵a的逆

3.3.5 爱因斯坦的求和约定

对张量的某些维度进行求和操作，也被称为张量的缩并。
例如，两个一维张量的缩并相当于两个矢量的点乘，一个
二阶张量的两个维度之间的缩并相当于矩阵的求迹。
所谓的爱因斯坦求和约定，是一种用于化简数学符号的约定。
这个约定是指在等号的同侧，相同的指标代表求和；在等号
的异侧相同的指标用于标记不同的等式，不作求和处理。