Nx examples

Mix.install([
  {:axon, "~> 0.6"},
  {:exla, "~> 0.7"},
  {:nx, "~> 0.7"}
])

Tensor examples

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

Nx.tensor([[1, 2], [3, 4], [5, 6]], type: {:f, 32}, names: [:x, :y])

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

c = Nx.multiply(a, b)

Nx.sum(c)

tensor = Nx.iota({3, 1000, 1000})

Nx.sum(tensor)

Defn examples

Simple module uses defn macro

defmodule MyModule do
  import Nx.Defn

  defn adds(x, n) do
    (x + n) |> print_expr()
  end

  defn adds2(x, n) do
    Nx.add(x, n) |> print_expr()
  end
end

a = Nx.tensor([[1, 2, 3], [4, 5, 6]])
MyModule.adds(a, 3)

MyModule.adds2(a, 5)

defmodule Softmax do
  import Nx.Defn

  defn softmax(n) do
    (Nx.exp(n) / Nx.sum(Nx.exp(n)))
    |> print_expr()
  end
end

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

Softmax.softmax(a)

Vectorization example

defmodule MyModule2 do
  import Nx.Defn

  defn normalize(t) do
    case Nx.shape(t) do
      {_} -> :ok
      _ -> raise "invalid shape"
    end

    min = Nx.reduce_min(t)
    zero_min = Nx.subtract(t, min)
    Nx.divide(zero_min, Nx.reduce_max(zero_min))
  end
end

MyModule2.normalize(Nx.tensor([3, 2, 6]))