Elixir Nx の基礎

import IEx.Helpers

Mix.install(
  [
    {:nx, "~> 0.4.0"},
    {:exla, "~> 0.4.0"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

概要

• Elixir Nx の基本的な概念や動作を学ぶ

資料

• https://hexdocs.pm/nx/Nx.html
• Elixir Nx の基礎 (Livebook)
• 深掘り Nx シリーズ ~ tensor を知る ① ~
• 深掘り Nx シリーズ ~ tensor を知る ② 算術計算 ~
• 深掘り Nx シリーズ ~ tensor を知る ③ 算術計算その 2 行列の積 ~

Nx

• 効率よく数値計算を行うためのライブラリ
• テンソル 上で計算を行う
• 3 つの能力
  1. テンソル
     • 型を持った多次元データ
     • 次元に名前を与えることができる
  2. 数値関数の定義 (defn)
     • カスタムコードを関数化できる
     • テンソルを扱える
  3. 自動微分 (Automatic differentiation)
     • autograd または autodiff として知られている
     • 共通の数値計算のシナリオをサポート
       • 機械学習
       • シュミレーション
       • 曲線回帰
       • 確率モデル
       • など
     • 機械学習やディープラーニングにおいて勾配計算を行う

Nx のテンソル

• 型を持った多次元データ
• 何重にも入れ子になった配列のイメージ
• 物理学で言う何らかの物理量を表すテンソルとは無関係
• 次元という言葉を使う

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

テンソルの生成

Nx.tensor/2

h(Nx.tensor())

0 次元テンソル

Nx.tensor(0)

1 次元テンソル

import Nx, only: :sigils
~V[1 2 3]f32

t1 = Nx.tensor([1.0, 2.0, 3.0], names: [:x])

t1[1]

t1[x: 1]

t1[x: 0..1]

# 行ベクトル
row = Nx.tensor([1, 2, 3])

# 列ベクトル
column = Nx.tensor([[1], [3], [8]])

2 次元テンソル

import Nx, only: :sigils

~M'''
1 2 3
4 5 6
'''s32

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

t2[0][2]

t2[x: 0][y: 2]

t2[x: 0..1][y: 1..2]

3 次元テンソル

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]
  )

t3[-1][-1][-1]

t3[batch: 1..-1//1]

Nx.shape/1 and Nx.reshape/2

Nx.shape(t3)

Nx.reshape(t3, {6, 3}, names: [:batches, :values])

Nx.to_binary/2 and Nx.from_binary/3

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

Nx.from_binary(<<0, 1, 2>>, :u8)

Access syntax (slicing)

t =
  Nx.tensor([
    [1, 2],
    [3, 4],
    [5, 6],
    [7, 8]
  ])

# Drop the first "row"
t[1..-1//1]

t =
  Nx.tensor([
    [1, 2],
    [3, 4],
    [5, 6],
    [7, 8]
  ])

# Drop the first "row" twice
t[1..-1//1][1..-1//1]

t =
  Nx.tensor([
    [1, 2],
    [3, 4],
    [5, 6],
    [7, 8]
  ])

# Drop the first "row" and the first "column"
t[[1..-1//1, 1..-1//1]]

t =
  Nx.tensor([
    [1, 2],
    [3, 4],
    [5, 6],
    [7, 8]
  ])

# Drop only the first "column"
t[[.., 1..-1//1]]

t =
  Nx.tensor([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9],
    [10, 11, 12]
  ])

t[[1..2, 2]]

テンソル関数

• 任意の shape のテンソルを引数にとれる関数

Nx.add(1, 2)

Nx.sum

t3[0] |> Nx.sum() |> dbg()

:ok

t2 |> Nx.sum(axes: [:x]) |> dbg()

:ok

t2 |> Nx.sum(axes: [:y]) |> dbg()

:ok

t2 |> Nx.sum(axes: [:x, :y]) |> dbg()

:ok

Nx.add

Nx.add(
  Nx.tensor([1, 2, 3]),
  Nx.tensor([10, 20, 30])
)

Nx.subtract

Nx.subtract(
  Nx.tensor([[1, 2], [3, 4]]),
  Nx.tensor([[5, 6], [7, 8]])
)

ブロードキャスト

# 1 を {2, 2} テンソルに変換
Nx.broadcast(1, {2, 2})

Nx.tensor([[1], [2]]) |> Nx.broadcast({1, 2, 4}) |> dbg()

:ok

数値関数の定義 (defn)

h(Nx.Defn)

自動微分 (autograd)

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

h(Nx.Defn.grad())

# sinの微分はcosなので、0 地点の勾配は 1 になる
fun = Nx.Defn.grad(&Nx.sin(&1))

Nx.tensor(0) |> fun.() |> dbg()

:ok

テンソルの軸 (Axis)

• テンソルの軸はテンソルの多次元配列の深さを示す

Nx.iota

• shape を与えると、0 始まりの連番の値を持つテンソルを作ってくれる

Nx.iota({5})

Nx.iota({2, 3}, names: [:x, :y])

Nx.axis_index

• テンソルの軸を与えると、テンソルの index を返す
• -1 は一番最後の軸に相当

0 = Nx.iota({2, 3, 4}) |> Nx.axis_index(0)
1 = Nx.iota({2, 3, 4}) |> Nx.axis_index(1)
2 = Nx.iota({2, 3, 4}) |> Nx.axis_index(2)
2 = Nx.iota({2, 3, 4}) |> Nx.axis_index(-1)

:ok

0 = Nx.iota({2, 3, 4}, names: [:x, :y, :z]) |> Nx.axis_index(:x)
1 = Nx.iota({2, 3, 4}, names: [:x, :y, :z]) |> Nx.axis_index(:y)
2 = Nx.iota({2, 3, 4}, names: [:x, :y, :z]) |> Nx.axis_index(:z)

:ok

Nx.axis_size

2 = Nx.iota({2, 3, 4}) |> Nx.axis_size(0)
3 = Nx.iota({2, 3, 4}) |> Nx.axis_size(1)
4 = Nx.iota({2, 3, 4}) |> Nx.axis_size(2)

:ok

2 = Nx.iota({2, 3, 4}, names: [:x, :y, :z]) |> Nx.axis_size(:x)
3 = Nx.iota({2, 3, 4}, names: [:x, :y, :z]) |> Nx.axis_size(:y)
4 = Nx.iota({2, 3, 4}, names: [:x, :y, :z]) |> Nx.axis_size(:z)

:ok

new_axis

Nx.tensor([[1, 2, 3], [4, 5, 6]]) |> Nx.new_axis(0, :new)

Nx.tensor([[1, 2, 3], [4, 5, 6]]) |> Nx.new_axis(1, :new)

Nx.tensor([[1, 2, 3], [4, 5, 6]]) |> Nx.new_axis(2, :new)

Nx.tensor([[1, 2, 3], [4, 5, 6]]) |> Nx.new_axis(-1, :new)

one-hot エンコーディング

• 例えば 3 というスカラー値を、[0, 0, 0, 1, 0, 0, 0, 0, 0, 0] という 0 番始まりの 3 番目だけ 1 で他は全て 0 のベクトルに変換すること

Nx.tensor(Enum.to_list(1..4))
|> Nx.new_axis(-1)
|> Nx.equal(Nx.tensor(Enum.to_list(0..9)))
|> dbg()

:ok