Hooks on EXLA backend

I’m trying to use hooks to debug my defn functions, but it works only on Default backend and doesn’t work on EXLA.
And I see the tests for that functionality on Nx/EXLA repo, so it should work. May be am doing something wrong, or forgot something?

Mix.install(
  [
    {:nx, "~> 0.3"},
    {:exla, "~> 0.3"},
  ],
  config: [
    nx: [
      default_backend: EXLA.Backend,
      default_defn_options: [compiler: EXLA]
    ]
  ]
)
defmodule Hooks do
  import Nx.Defn

  defn add_and_mult(a, b) do
    add = hook(a + b, fn tensor -> IO.inspect({:add, tensor}) end)
    mult = hook(a * b, fn tensor -> IO.inspect({:mult, tensor}) end)
    {add, mult}
  end
end

Hooks.add_and_mult(2, 3)

It shows only result.

1 Like

I found out what’s going on.
I’m using livebook and the output from EXLA backend goes directly to livebook console.

1 Like

Have you tried using ielixir?