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())