随着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]  ]>

加减运算

tensor 的相加和相减运算, 是每个雷同地位上的元素进行相加和相减. 例如 t + t 会失去:

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

形态不同的两个 tensor, 能不能进行加减呢? 答案是可能能够, 如果其中一个 tensor 能够被扩散(broadcast) 成另外一个 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 , 它能够扩大为 2x1, 即 t2 的形态:

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