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.