Nx/EXLA benchmarking

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

Nx.default_backend(EXLA.Backend)
Nx.Defn.global_default_options(compiler: EXLA)

🐈

defmodule Calc do
  import Nx.Defn

  defn cube_power(t) do
    Nx.power(t, 3)
  end

  defn cube_multiply(t) do
    t * t * t
  end
end

Benchmark

defmodule Benchmark do
  def run() do
    t = Nx.iota({1000, 1000})

    # Useful when comparing different numerical implementations
    Benchee.run(%{
      "cube_power" => fn -> Calc.cube_power(t) end,
      "cube_multiply" => fn -> Calc.cube_multiply(t) end
    })
  end
end

Benchmark.run()