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