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.