Chapt 2

Thinking in Tensors

  Mix.install([
      {:axon, "~> 0.6.1"},
      {:nx, "~> 0.7.3"},
      {:exla, "~> 0.7.3"},
      {:bumblebee, "~> 0.5.3"},
      {:explorer, "~> 0.9.1"},
      {:kino, "~> 0.13.2"},
      {:benchee, "~> 1.3"}
  ])

a = Nx.tensor([1, 2, 3])
b = Nx.tensor([[1, 2, 3], [4, 5, 6]])
c = Nx.tensor(1.0)
d = Nx.tensor([[[[[[1.0, 2]]]]]])
e = Nx.to_binary(b)
f = Nx.as_type(a, {:f, 32}) |> Nx.reshape({1, 3, 1})

for i <- [a, b, c, d, e, f], do: dbg(i)

defmodule TestingDefn do
  import Nx.Defn

  defn adds_one(t) do
    Nx.add(t, 1)
  end
end

TestingDefn.adds_one(b)

defmodule TestingSoftmax do
  import Nx.Defn

  defn soft_max(t), do: Nx.exp(t) / Nx.sum(Nx.exp(t))
end

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 w/ EXLA" => fn ->
      apply(EXLA.jit(&TestingSoftmax.soft_max/1), [tensor])
      end,
    "Regular ELixir" => fn ->
      TestingSoftmax.soft_max(tensor)
    end
  },
  time: 10
)