From 5e2e61bceb455d8f410c76f10a1665784b9ea654 Mon Sep 17 00:00:00 2001 From: Marcelo Reichert Date: Wed, 12 Aug 2026 15:36:12 -0300 Subject: [PATCH 1/2] Accept :hooks compiler opt in EMLX.__compile__/4 nx's `Nx.Defn.compile/3`/`jit/2` now always forward a `:hooks` option (defaulting to `%{}`) to the selected compiler, even when no hooks are used. EMLX's strict Keyword.validate! allowlist didn't include it, so any compile under a current nx (e.g. bench/svd_bench.exs, which pins nx main) failed immediately with "unknown keys [:hooks]". An empty hooks map is now accepted and ignored. A non-empty one raises a clear error instead: EMLX lowers :hook/:io_call expr nodes natively (each with its own inline default callback), but doesn't wire through the named-override map carried by this compiler opt. --- emlx/lib/emlx.ex | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/emlx/lib/emlx.ex b/emlx/lib/emlx.ex index 7aab32e..fe0f902 100644 --- a/emlx/lib/emlx.ex +++ b/emlx/lib/emlx.ex @@ -1741,7 +1741,16 @@ defmodule EMLX do # Known EMLX-specific compiler opts. `:command_queue` is injected by # `__partitions_options__/1` but may also be passed directly by callers # that manage their own queues (equivalent to a manual `with_queue`). - @valid_compiler_keys [:device, :max_concurrency, :command_queue] + # + # `:hooks` isn't EMLX-specific — `Nx.Defn.compile/3`/`jit/2` always forward + # it (defaulting to `%{}`) to whichever compiler is selected (see + # `Nx.Defn.prepare_options/1`), so every `Nx.Defn.Compiler` callback must + # accept it even if unused. `:hook`/`:io_call` *expr nodes* (each carrying + # their own inline default callback) do lower natively — see + # `EMLX.Native.Expr`'s `:hook`/`:io_call` clause — but the *named-override* + # map this option carries (swapping in a different callback per hook name + # at compile time) isn't wired through, so it's only accepted when empty. + @valid_compiler_keys [:device, :max_concurrency, :command_queue, :hooks] # Process-lifetime dispatch cache backing `dispatch_key/3` + # `get_or_compile_program/6` (see their docs) — a compiled program is keyed @@ -1772,6 +1781,19 @@ defmodule EMLX do @impl Nx.Defn.Compiler def __compile__(_key, vars, fun, opts) do Keyword.validate!(opts, @valid_compiler_keys) + + case Keyword.get(opts, :hooks, %{}) do + empty when empty == %{} -> + :ok + + hooks -> + raise ArgumentError, + "EMLX does not support the :hooks named-override map (got callbacks for " <> + "#{inspect(Map.keys(hooks))}) — :hook/:io_call expr nodes lower natively " <> + "with their own inline default callback, but swapping in a different " <> + "callback per hook name via compiler opts isn't wired through" + end + queue = Keyword.get(opts, :command_queue) device = Keyword.get(opts, :device, default_device()) From a6ae72b33eda0d19e5ff811028bc7e19019b3915 Mon Sep 17 00:00:00 2001 From: Marcelo Reichert Date: Wed, 12 Aug 2026 15:36:12 -0300 Subject: [PATCH 2/2] Accept :hooks compiler opt in EMLX.__compile__/4 nx's `Nx.Defn.compile/3`/`jit/2` now always forward a `:hooks` option (defaulting to `%{}`) to the selected compiler, even when no hooks are used. EMLX's strict Keyword.validate! allowlist didn't include it, so any compile under a current nx (e.g. bench/svd_bench.exs, which pins nx main) failed immediately with "unknown keys [:hooks]". An empty hooks map is now accepted and ignored. A non-empty one raises a clear error instead: EMLX lowers :hook/:io_call expr nodes natively (each with its own inline default callback), but doesn't wire through the named-override map carried by this compiler opt. --- emlx/lib/emlx.ex | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/emlx/lib/emlx.ex b/emlx/lib/emlx.ex index 7aab32e..e1a847c 100644 --- a/emlx/lib/emlx.ex +++ b/emlx/lib/emlx.ex @@ -1741,7 +1741,7 @@ defmodule EMLX do # Known EMLX-specific compiler opts. `:command_queue` is injected by # `__partitions_options__/1` but may also be passed directly by callers # that manage their own queues (equivalent to a manual `with_queue`). - @valid_compiler_keys [:device, :max_concurrency, :command_queue] + @valid_compiler_keys [:device, :max_concurrency, :command_queue, :hooks] # Process-lifetime dispatch cache backing `dispatch_key/3` + # `get_or_compile_program/6` (see their docs) — a compiled program is keyed @@ -1772,6 +1772,19 @@ defmodule EMLX do @impl Nx.Defn.Compiler def __compile__(_key, vars, fun, opts) do Keyword.validate!(opts, @valid_compiler_keys) + + case Keyword.get(opts, :hooks, %{}) do + empty when empty == %{} -> + :ok + + hooks -> + raise ArgumentError, + "EMLX does not support the :hooks named-override map (got callbacks for " <> + "#{inspect(Map.keys(hooks))}) — :hook/:io_call expr nodes lower natively " <> + "with their own inline default callback, but swapping in a different " <> + "callback per hook name via compiler opts isn't wired through" + end + queue = Keyword.get(opts, :command_queue) device = Keyword.get(opts, :device, default_device())