NX Practice
Mix.install ([
{:nx, "~> 0.5"},
{:exla, "~> 0.5"},
{:benchee, github: "bencheeorg/benchee", override: true}
])
Main
# Create Tensor
Nx.tensor([1, 2, 3])
# Create More Tensors
a = Nx.tensor([[1, 2, 3], [4, 5, 6]])
b = Nx.tensor(1.0)
c = Nx.tensor([[[[[[1.0, 2]]]]]])
dbg(a)
dbg(b)
dbg(c)
# Tensor Types Example
a = Nx.tensor([1, 2, 3])
b = Nx.tensor([1.0, 2.0, 3.0])
dbg(a)
dbg(b)
# Underflow Example
Nx.tensor(0.00000000000000000000000000000000000000000000000001)
Nx.tensor(1.0e-45, type: {:f, 64})
# Overflow Example
Nx.tensor(128, type: {:s, 8})
# Auto Type Assign Example
Nx.tensor([1.0, 2, 3])
# Example of Tensor's Shape
a = Nx.tensor([1, 2])
b = Nx.tensor([[1, 2], [3, 4]])
c = Nx.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
dbg(a)
dbg(b)
dbg(c)
# Example of a Scalar
Nx.tensor(10)
# Example of Named Tensors
Nx.tensor([[1, 2, 3], [4, 5, 6]], names: [:x, :y])
# Example of Binary Representation
a = Nx.tensor([[1, 2, 3], [4, 5, 6]])
Nx.to_binary(a)
# It's more performant to turn your data into binary representation:
<<1::64-signed-native, 2::64-signed-native, 3::64-signed-native>>
|> Nx.from_binary({:s, 64})
# Changing the Shape of a Tensor using reshape:
<<1::64-signed-native, 2::64-signed-native, 3::64-signed-native>>
|> Nx.from_binary({:s, 64})
|> Nx.reshape({1, 3})
# Shaping a Tensor
a = Nx.tensor([1, 2, 3])
a
|> Nx.as_type({:f, 32})
|> Nx.reshape({1, 3, 1})
# Bitcast Example
Nx.bitcast(a, {:f, 64})
# Element-wise Unary Operations
a = Nx.tensor([[[-1, -2, -3], [-4, -5, -6]], [[1, 2, 3], [4, 5, 6]]])
Nx.abs(a)
# Element-wise Binary Operations
a = Nx.tensor([[1, 2, 3], [4, 5, 6]])
b = Nx.tensor([[6, 7, 8], [9, 10, 11]])
Nx.add(a, b)
Nx.multiply(a, b)
Nx.add(5, Nx.tensor([1, 2, 3]))
Nx.add(Nx.tensor([1, 2, 3]), Nx.tensor([[4, 5, 6], [7, 8, 9]]))
# Reductions
revs = Nx.tensor([85, 76, 42, 34, 46, 23, 52, 99, 22, 32, 85, 51])
Nx.sum(revs)
revs = Nx.tensor(
[
[21, 64, 86, 26, 74, 81, 38, 79, 70, 48, 85, 33],
[64, 82, 48, 39, 70, 71, 81, 53, 50, 67, 36, 50],
[68, 74, 39, 78, 95, 62, 53, 21, 43, 59, 51, 88],
[47, 74, 97, 51, 98, 47, 61, 36, 83, 55, 74, 43]
], names: [:year, :month])
Nx.sum(revs, axes: [:year])
Nx.sum(revs, axes: [:month])
defmodule MyModule do
def adds_one(x) do
Nx.add(x, 1)
end
end
defmodule DefnModule do
import Nx.Defn
defn adds_one(x) do
Nx.add(x, 1) |> print_expr()
end
end
DefnModule.adds_one(Nx.tensor([1, 2, 3]))
# Defn Acceleration Comparison
defmodule Softmax do
import Nx.Defn
defn softmax(n), do: Nx.exp(n) / Nx.sum(Nx.exp(n))
end
key = Nx.Random.key(42)
{tensor, _key} = Nx.Random.uniform(key, shape: {1_000_000})
Benchee.run(
%{
"JIT with EXLA" => fn ->
apply(EXLA.jit(&Softmax.softmax/1), [tensor])
end,
"Regular Elixir" => fn ->
Softmax.softmax(tensor)
end
},
time: 10
)