Chapter 2

Mix.install([
  {:nx, "~> 0.9"},
  {:exla, "~> 0.9"},
  {:benchee, "~> 1.3"}
])

Understanding Nx Tensors

Nx.tensor([[[[1,2,3],[1,2,3]],[[1,2,3],[1,2,3]]],[[[1,2,3],[1,2,3]],[[1,2,3],[1,2,3]]]])

a = Nx.tensor([[1,2,3],[4,5,6]])
b = Nx.tensor(2.0)
c = Nx.tensor([[[[[[[[[[4.5]]]]]]]]]])
dbg(a)
dbg(b)
dbg(c)
nil

a = Nx.tensor([1, 2, 3])
b = Nx.tensor([1.0, 2.0, 3.0])
dbg(a)
dbg(b)
nil

dbg(Nx.tensor(0.0000000000000000000000000000000000000000000001))
dbg(Nx.tensor(0.0000000000000000000000000000000000000000000001, type: {:f, 64}))
nil

# I guess I will spend my time debugging this kind of overflow.
Nx.tensor(128, type: {:s, 8})

Nx.tensor([1.0, 2, 3])

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

Nx.tensor([[1,2,3],[4,5,6]], names: [:x, :y])
|> Nx.to_binary()

<<1::64-signed-native, 2::64-signed-native, 3::64-signed-native>>
|> Nx.from_binary({:s, 64})
|> Nx.reshape({1,1,1,3})

Element-wise Unary Operations

a = Nx.tensor([1,2,3])

a
|> Nx.as_type({:f, 32})
|> Nx.reshape({1,3,1})

# Nx.bitcast(a, {:f, 32})

# Elixir way (this will be problematic in higher rank tensors)
a = [-1, -2, -3, 0, 1, 2, 3]
Enum.map(a, &abs/1)

# Nx way
a = Nx.tensor([[[-1,-2,-3],[-4,-5,-6]], [[1,2,3],[4,5,6]]])
Nx.abs(a)

Element-wise Binary Operations

# Elixir way
a = [1,2,3]
b = [4,5,6]

Enum.zip_with(a, b, fn x, y -> x + y end)

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)

# Broadcasting
dbg(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])
dbg(Nx.sum(revs))
dbg(Nx.mean(revs))
nil

revs =
  Nx.tensor([
    [85, 76, 42, 34, 46, 23, 52, 99, 22, 32, 85, 51],
    [64, 82, 48, 39, 70, 71, 81, 53, 50, 67, 36, 50],
    [68, 76, 42, 34, 46, 23, 52, 99, 22, 32, 51, 88],
    [85, 42, 42, 12, 46, 63, 23, 10, 32, 54, 81, 51]
  ], names: [:year, :month])

dbg(Nx.sum(revs, axes: [:year]))
dbg(Nx.sum(revs, axes: [:month]))
dbg(Nx.sum(revs, axes: [:year, :month]))
nil

Going from def to defn

defmodule MyModule do
  import Nx.Defn
  
  defn adds_one(x) do
    Nx.add(x, 1) |> print_expr()
  end
end

MyModule.adds_one(Nx.tensor([1,2,3]))

defmodule Softmax do
  import Nx.Defn

  defn softmax(n), do: Nx.exp(n) / Nx.sum(Nx.exp(n))
end

# Set to always use EXLA
#Nx.Defn.global_default_options(compiler: EXLA)

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
)