Elixir Nx の基礎 (Livebook)

Elixir **Nx** の基本的な概念や動作を、以下のドキュメントに沿って確認しました。Livebook 上で動作させています。 
[Introduction to Nx 公式ドキュメント](https://hexdocs.pm/nx/0.4.0/intro-to-nx.html)



**関連記事**
[Elixir Nx で学ぶ機械学習 - Axon無し (Livebook) - Qiita](https://qiita.com/sand/items/1ad04bafd4ef3d91a675)
[Elixir Axon で学ぶ機械学習 (Livebook) - Qiita](https://qiita.com/sand/items/390b14a8665be2c0baf9)
[Elixir Nx の基礎 (Livebook) - Qiita](https://qiita.com/sand/items/9a47cfe0c95386417224)
[Livebook 事始め (Elixir) - Qiita](https://qiita.com/sand/items/6753c1782774b45daf6b)

# 1. Nx 概要

まず Livebook 上で必要なインストールを行っておきます。
```elixir:
Mix.install([
  {:nx, "~> 0.2"}
])
```


**Nx** は Elixir が効率よく数値計算を行うためのライブラリです。それは **型を持った多次元データ**  である **テンソル** 上で計算を行います。**Nx** は基本的に次の **3** つの能力を持っています。

* **Nx** において **テンソル** は **型を持ったデータ** であり、**多次元**であり、しかも次元には **名前** を与えることができる。
* 数値関数の定義 (**defn**) で カスタムコードを関数化できます。それは **tensor-aware** オペレーションであり、関数です。
* **自動微分** (Automatic differentiation) は、 **autograd** または **autodiff** として知られており、共通の数値計算のシナリオをサポートしています。**機械学習** やシュミレーション, 曲線回帰、確率モデルなどです。

::: note warn
テンソルの定義は、型を持った多次元データである、ということです。単に、何重にも入れ子になった配列のイメージです。物理学で言う何らかの物理量を表すテンソルとは無関係です。  
:::

::: note info
自動微分は機械学習やディープラーニングにおいて勾配計算を行う、最重要な概念です。
:::

何故、**微分** が必要になるかについては、興味があれば以下の過去記事を参照してください。
[【機械学習】誤差逆伝播法のコンパクトな説明](https://qiita.com/sand/items/2d783a12c575fb949c6e)

Nx テンソルは以下の型を持ちます。
* **unsigned integers**  - u8, u16, u32, u64
* **signed integers** - s8, s16, s32, s64
* **floats** - f32, f64
* **brain floats** - bf16
* **complex** - c64, c128


# 2. Nx のテンソル
Nx のテンソルは、何ら物理量とは対応していないといいましたが、テンソルという言葉や概念は物理学由来だと思います。物理学ではスカラーは零階のテンソル、ベクトルは一階のテンソル、二階のテンソルと呼ぶらしいですが、ここでは次元という言葉を使っています。

* 0次元テンソル - スカラー
* 1次元テンソル - ベクトル
* 2次元テンソル - 行列
* 3次元テンソル - 行列の配列
* ...
* n次元テンソル - (n-1)次元テンソルの配列 


## 2-1. テンソルの作成

Elixir リストからテンソルを作成することができます。以下のヘルプに詳細な例がありますが、いくつか見ていきましょう。

```elixir:
import IEx.Helpers
h Nx.tensor
```

### 0次元テンソル 

```
    Nx.tensor(0)
    #Nx.Tensor<
      s64
      0
    >
```


### 1次元テンソル

```
    Nx.tensor([1, 2, 3])
    #Nx.Tensor<
      s64[3]
      [1, 2, 3]
    >

    Nx.tensor([1.2, 2.3, 3.4, 4.5])
    #Nx.Tensor<
      f32[4]
      [1.2000000476837158, 2.299999952316284, 3.4000000953674316, 4.5]
    >
```

型を明示的に指定できる。型指定で最大値を超えた場合はオーバフローとなる。
2進数(1,000,000 300) = 10進数(256)なので、 8ビットでオーバーフローを考慮すると、300 -256 = 44 になる。自動的に以下のような変換が行われる。


```
    Nx.tensor([300, 301, 302], type: :s8)
    #Nx.Tensor<
      s8[3]
      [44, 45, 46]
    >
```

### 2次元テンソル

```
    Nx.tensor([[1, 2, 3], [4, 5, 6]])
    #Nx.Tensor<
      s64[2][3]
      [
        [1, 2, 3],
        [4, 5, 6]
      ]
    >
    
```

これは以下のような行列に相当します。

```math
\begin{pmatrix} 1 & 2 & 3\\ 4 & 5 & 6 \end{pmatrix}
```


### 3次元テンソル


```
Nx.tensor([[[1, 2], [3, 4], [5, 6]], [[-1, -2], [-3, -4], [-5, -6]]])
#Nx.Tensor<
  s64[2][3][2]
  [
    [
      [1, 2],
      [3, 4],
      [5, 6]
    ],
    [
      [-1, -2],
      [-3, -4],
      [-5, -6]
    ]
  ]
>
```


### 次元軸に名前を付ける

テンソル次元に名前を付けることができます。コードの可読性を高めてくれます。**names は atom のリストですが、テンソルの多次元配列の入れ子の浅い軸の順に並びます。**

```
    t1 = Nx.tensor([1, 2, 3], names: [:x])
    #Nx.Tensor<
      s64[x: 3]
      [1, 2, 3]
    >

    t2 = Nx.tensor([[1, 2, 3], [4, 5, 6]], names: [:x, :y])
    #Nx.Tensor<
      s64[x: 2][y: 3]
      [
        [1, 2, 3],
        [4, 5, 6]
      ]
    >

    t3 = Nx.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
      [[-1, -2, -3], [-4, -5, -6], [-7, -8, -9]]], names: [:batch, :height, :width])
    #Nx.Tensor<
      s64[batch: 2][height: 3][width: 3]
      [
        [
          [1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]
        ],
        [
          [-1, -2, -3],
          [-4, -5, -6],
          [-7, -8, -9]
        ]
      ]
    >
```

名前でテンソルの要素にアクセスしてみます。あわせて **Nx.shape/1** でテンソルの shape を取得し、**Nx.reshape/2** で shape を変更します。
```
t1[x: 1]
#Nx.Tensor<
  s64
  2
>

t2[x: 1]
#Nx.Tensor<
  s64[y: 3]
  [4, 5, 6]
>

t2[y: 1]
#Nx.Tensor<
  s64[x: 2]
  [2, 5]
>

t2[x: 0, y: 1]
#Nx.Tensor<
  s64
  2
>

t3[height: 1]
#Nx.Tensor<
  s64[batch: 2][width: 3]
  [
    [4, 5, 6],
    [-4, -5, -6]
  ]
>

t3[width: 2]
#Nx.Tensor<
  s64[batch: 2][height: 3]
  [
    [3, 6, 9],
    [-3, -6, -9]
  ]
>


Nx.shape(t1)
{3}

Nx.shape(t2)
{2, 3}

Nx.shape(t3)
{2, 3, 3}

Nx.reshape(t3, {6, 3}, names: [:batches, :values])
#Nx.Tensor<
  s64[batches: 6][values: 3]
  [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9],
    [-1, -2, -3],
    [-4, -5, -6],
    [-7, -8, -9]
  ]
>
```

### バイナリ変換


```
Nx.tensor([[1, 2], [3, 4]], type: :u8) |> Nx.to_binary()
<<1, 2, 3, 4>>

Nx.from_binary(<<0, 1, 2>>, :u8)
#Nx.Tensor<
  u8[3]
  [0, 1, 2]
>
```

## 2-2. テンソル関数

テンソル関数とは任意の shape のテンソルを引数にとれる関数です。
```
Nx.sum(t3[0])
#Nx.Tensor<
  s64
  45
>

Nx.sum(t2)
#Nx.Tensor<
  s64
  21
>

Nx.sum(t2, axes: [:x])
#Nx.Tensor<
  s64[y: 3]
  [5, 7, 9]
>
```

```
tensor1 = Nx.tensor([[1, 2], [3, 4]])
tensor2 = Nx.tensor([[5, 6], [7, 8]])
Nx.subtract(tensor2, tensor1)
#Nx.Tensor<
  s64[2][2]
  [
    [4, 4],
    [4, 4]
  ]
>

```



## 2-3. ブロードキャスト

テンソルでの計算ですが、例えば以下のような引き算は、
```math
\begin{pmatrix} 1 & 2 & 3\\ 4 & 5 & 6 \end{pmatrix} - 1 = \begin{pmatrix} 0 & 1 & 2\\ 3 & 4 & 5 \end{pmatrix}
```

数学的な意味では以下の引き算に置き換えられると、便利です。
```math
\begin{pmatrix} 1 & 2 & 3\\ 4 & 5 & 6 \end{pmatrix} - \begin{pmatrix} 1 & 1 & 1\\ 1 & 1 & 1 \end{pmatrix} = \begin{pmatrix} 0 & 1 & 2\\ 3 & 4 & 5 \end{pmatrix}
```

**Nx.broadcast/2** を使えば、1 を {2, 2} テンソルに変換してくれます。
```
Nx.broadcast(1, {2, 2})
#Nx.Tensor<
  s64[2][2]
  [
    [1, 1],
    [1, 1]
  ]
>
```

**Nx.broadcast(1, tensor)** は 1 を tensor の shape にブロードキャストしてくれます。また例えば **Nx.subtract(tensor, 1)** も2つの引数のテンソルの shape を互換性のあるものへと自動的にブロードキャストしてくれます。

```
[[1], [2]] |> Nx.tensor() |> Nx.broadcast({1, 2, 4})
#Nx.Tensor<
  s64[1][2][4]
  [
    [
      [1, 1, 1, 1],
      [2, 2, 2, 2]
    ]
  ]
>
```

以下のヘルプで多くの example が表示できます。

```
h Nx.broadcast
```



# 3. 数値関数の定義 (defn)

**defn** を使って好きなテンソル関数を作り出すことができます。引数として、スカラーやベクトル、行列などのテンソルを受けることができます。

引き算と足し算の関数を定義します。
```elixir:
defmodule TensorMath do
  import Nx.Defn

  defn hiku(a, b) do
    a - b
  end

  defn tasu(a, b) do
    a + b
  end
end
```


行列テンソル同士の引き算です。
```
tensor1 = Nx.tensor([[1, 1], [1, 1]])
tensor2 = Nx.tensor([[5, 6], [7, 8]])
TensorMath.hiku(tensor2, tensor1)
#Nx.Tensor<
  s64[2][2]
  [
    [4, 5],
    [6, 7]
  ]
>
```

tensor1 をスカラー1に変えても同じ結果です。
```
TensorMath.hiku(tensor2, 1)
#Nx.Tensor<
  s64[2][2]
  [
    [4, 5],
    [6, 7]
  ]
>
```

足し算についても同じ結果が得られます。
```
TensorMath.tasu(tensor2, tensor1)
#Nx.Tensor<
  s64[2][2]
  [
    [6, 7],
    [8, 9]
  ]
>

TensorMath.tasu(tensor2, 1)
#Nx.Tensor<
  s64[2][2]
  [
    [6, 7],
    [8, 9]
  ]
>
```


# 4. 自動微分 (autograd)

```
h Nx.Defn.grad
```

Nx.Defn.grad(fun) は無名関数を引数として、新たな無名関数を作成し返してくれます。その無名関数は与えられた地点の勾配を求めてくれます。

sinの微分はcosですので、0 地点の勾配は 1 になります。
```
    fun = Nx.Defn.grad(fn x -> Nx.sin(x) end)
    fun.(Nx.tensor(0))
    #Nx.Tensor<
      f32
      1.0
    >
```


以下のような f(x) 関数の導関数は g(x) となります。
```math:
f(x) = 3x^{2} + 2x + 1\\
g(x) = 6x + 2
```

この事実をもとに Nx の自動微分を行ってみます。
地点 2 の勾配を求めてみます。

```
defmodule Funs do
  import Nx.Defn

  defn poly(x) do
    3 * Nx.power(x, 2) + 2 * x + 1
  end

  defn poly_slope_at(x) do
    grad(&poly/1).(x)
  end
end

Funs.poly_slope_at(2)
#Nx.Tensor<
  f32
  14.0
>
```

g(2) = 14 の結果とあっています。



# 5. テンソルの軸 (Axis)

テンソルの軸は少しわかりにくい概念なので、もう少し述べておきたいと思います。
便利な関数なのですが、**Nx.iota()** に shape を与えると、0　始まりの連番の値を持つテンソルを作ってくれます。

```
Nx.iota({5})
#Nx.Tensor<
  s64[5]
  [0, 1, 2, 3, 4]
>

Nx.iota({3, 2, 3}, names: [:batch, :height, :width])
#Nx.Tensor<
  s64[batch: 3][height: 2][width: 3]
  [
    [
      [0, 1, 2],
      [3, 4, 5]
    ],
    [
      [6, 7, 8],
      [9, 10, 11]
    ],
    [
      [12, 13, 14],
      [15, 16, 17]
    ]
  ]
>
```



テンソルの軸はテンソルの多次元配列の深さを示すものですが、**axis_index(tensor, axis)** はテンソルの軸を与えることによって、テンソルの index を返してくれるものです。
-1 は一番最後の軸に相当します。
```
Nx.axis_index(Nx.iota({100, 10, 20}), 0)
0

Nx.axis_index(Nx.iota({100, 10, 20}), -1)
2

Nx.axis_index(Nx.iota({100, 10, 20}, names: [:batch, :x, :y]), :x)
1
```



**axis_size(tensor, axis)** は与えられた軸のサイズを返します。
```
Nx.axis_size(Nx.iota({100, 10, 20}), 0)
100

Nx.axis_size(Nx.iota({100, 10, 20}, names: [:batch, :x, :y]), :y)
20
```


**new_axis(tensor, axis, name \\ nil)** はサイズ1 の軸を追加します。
```
t = Nx.tensor([[1, 2, 3], [4, 5, 6]])
#Nx.Tensor<
  s64[2][3]
  [
    [1, 2, 3],
    [4, 5, 6]
  ]
>

Nx.new_axis(t, 0, :new)
#Nx.Tensor<
  s64[new: 1][2][3]
  [
    [
      [1, 2, 3],
      [4, 5, 6]
    ]
  ]
>

Nx.new_axis(t, 1, :new)
#Nx.Tensor<
  s64[2][new: 1][3]
  [
    [
      [1, 2, 3]
    ],
    [
      [4, 5, 6]
    ]
  ]
>

Nx.new_axis(t, 2, :new)
#Nx.Tensor<
  s64[2][3][new: 1]
  [
    [
      [1],
      [2],
      [3]
    ],
    [
      [4],
      [5],
      [6]
    ]
  ]
>

Nx.new_axis(t, -1, :new)
#Nx.Tensor<
  s64[2][3][new: 1]
  [
    [
      [1],
      [2],
      [3]
    ],
    [
      [4],
      [5],
      [6]
    ]
  ]
>

```



## （例）one-hot エンコーディング 

one-hot エンコーディング とは、例えば 3 というスカラー値を、[0, 0, 0, 1, 0, 0, 0, 0, 0, 0] という 3番目だけ 1 で他は全て 0 のベクトルに変換することです。0 番始まりの 3 番目です。

わかり易いように段階を踏みます。


```
left =
  1..4
  |> Enum.to_list
  |> Nx.tensor
  |> Nx.new_axis(-1)
#Nx.Tensor<
  s64[4][1]
  [
    [1],
    [2],
    [3],
    [4]
  ]
>

right = Nx.tensor(Enum.to_list(0..9))
#Nx.Tensor<
  s64[10]
  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
>
```

**Nx.equal/2** は2つの tensor を比べて、等しい箇所を１に、そうでない個所を０にします。

```
Nx.equal(left, right)
#Nx.Tensor<
  u8[4][10]
  [
    [0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 1, 0, 0, 0, 0, 0]
  ]
>
```

（補足）内部では以下のようなブロードキャストが行われた後で、Nx.equal で要素ごとの比較が行われているものと推測されます。

```
left = 
  [
    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
    [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
    [3, 3, 3, 3, 3, 3, 3, 3, 3. 3],
    [4, 4, 4, 4, 4, 4, 4, 4, 4, 4]
  ]

right =
  [
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],  
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],  
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],  
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
  ]
```


今回は以上です。