开始学习多维类型数列(tensor)的计算

随着GPU 和 TPU 等异构芯片被广泛使用, 俺也逐渐不满足与只使用 CPU 进行计算. 为了刚上时代发展的浪潮, 俺决定学习一下如何写出可以在 GPU 上运行的程序. 最基础的是, GPU 的特性是多核并行计算, 所以需要使用特殊的数据结构, 例如多维数据结构(multidimensional data structures), 又被称为 tensor, 也有翻译为"张量".

恰好最近有一个 elixir 的 tensor 计算库发布, 名为 "Nx". 这里就使用它来学习一下 tensor 的基本计算.

新建一个 tensor

使用如下函数我们就新建了一个 2x2 的多维数组, 其中元素的类型是 64 位有符号整数(s64).

iex> t = Nx.tensor([[1, 2], [3, 4]])
#Nx.Tensor<
  s64[2][2]
  [
    [1, 2],
    [3, 4]
  ]
>

形状(shape),维度(rank)和体积(size)

每个 tensor 都有自己的形状, 可以用元组来表示, 例如一个 2 行 2 列的 tensor 的形状就是 {2, 2}, 维度是 2. 这样我们可以表示更多维度的 tensor, 例如一个形状为 {1, 2, 3, 4} 的 tensor, 其维度(rank)是 4.

标量的 rank 是 0, 形状是 {}.

和二维和三维集合一样, 体积是指一个 tensor 里元素的数量, 等于 shape 里所有边的乘积.

加减乘除运算

tensor 的加减乘除运算, 是每个相同位置上的元素(Element-wise)进行加减乘除. 例如 t + t 会得到:

#Nx.Tensor<
  s64[2][2]
  [
    [2, 4],
    [6, 8]
  ]
>

形状不同的两个 tensor, 能不能进行加减乘除呢? 答案是可能可以, 即把两个 tensor 的形状变为相同的, 再加减乘除.

iex(1)> t1 = Nx.tensor([1])
#Nx.Tensor<
  s64[1]
  [1]
>
iex(2)> t2 = Nx.tensor([[1], [2]])
#Nx.Tensor<
  s64[2][1]
  [
    [1],
    [2]
  ]
>
iex(3)> Nx.add(t1, t2)
#Nx.Tensor<
  s64[2][1]
  [
    [2],
    [3]
  ]
>

这里 t1 的形状是 1x1 , 它可以扩展(broadcast)为 2x1, 即 t2 的形状. 像一个卷轴, 卷起来的时候是一维的, 展开后就变成二维的.

iex> Nx.broadcast(t1, {2, 1})
#Nx.Tensor<
  s64[2][1]
  [
    [1],
    [1]
  ]
>

类似的以 Element-wise 的形式进行的运算还有 power(乘方) 和 remainder(求余).