M22: compile-time debug flags

Two opt-in Application.compile_env flags re-enable runtime assertions
on hot paths that GPU backends skip by default. Both default false
with zero runtime cost when off — the gated branches are dead-code
eliminated by the Elixir compiler and no Emily.Backend.DebugHelpers
reference reaches the compiled BEAM.

- :debug_bounds_check — raises on OOB / negative indices in gather,
  take, take_along_axis, indexed_add, indexed_put. Closes the silent-
  NaN-from-OOB-gather class of bug that surfaced in DistilBERT-QA
  conformance (vocab-30522 tokenizer vs. tiny-random embedding).
- :debug_detect_nan_inf — scans matmul, fast_rms_norm, fast_layer_norm,
  and both fast_scaled_dot_product_attention variants for NaN/Inf so
  training-time numerics failures surface at the producing op.

The Emily.Backend.DebugHelpers module holds the assertion bodies and
is always compiled; Emily.Backend only references it from inside
`if @flag` gates, so when the flags are off the entire module is
unreachable from Emily.Backend's bytecode. A zero-cost verification
test inspects :beam_disasm.file/1 output and asserts no reference
survives under the default-off build.

Emily's own config/test.exs keeps both prod flags false so mix test
doesn't pay GPU-sync cost on every run; fixture-only flags drive the
gate→helper composition tests in test/support/debug_fixture.ex.

Standalone softmax has no Backend callback in Emily (it's Axon-composed
from exp/sum/divide); the fused SDPA softmax is covered via the two
SDPA callback paths.

Closes #32.

Author: ausimian

README.md

@@ -113,6 +113,33 @@ config :emily, :warn_on_fallback, true
 The warning is off by default so library consumers and CI logs stay
 quiet. The telemetry event fires regardless.
 
+## Debug assertions
+
+Two compile-time flags re-enable runtime checks that MLX (and every
+other GPU backend) skips by default. Both are off by default with
+zero runtime cost when off — the guarded branches are dead-code
+eliminated by the Elixir compiler.
+
+```elixir
+# config/dev.exs
+config :emily,
+  debug_bounds_check: true,
+  debug_detect_nan_inf: true
+```
+
+- `:debug_bounds_check` — raises on out-of-range / negative indices
+  in `gather` / `take` / `take_along_axis` / `indexed_add` /
+  `indexed_put`. Catches the silent-`NaN`-from-OOB-gather class of
+  bug (e.g. a vocab-30522 tokenizer paired with a tiny-random model
+  whose embedding table is smaller).
+- `:debug_detect_nan_inf` — scans results of `matmul`, the fused
+  `layer_norm` / `rms_norm`, and both fused SDPA variants. Surfaces
+  numerics failures at the producing op rather than downstream.
+
+Each check is a per-op MLX reduction plus a scalar readback — a
+worker sync that breaks lazy-graph fusion. Leave off in release
+builds. See the `Emily` moduledoc for the full opt-in snippet.
+
 ## Milestones shipped
 
 - **M0** — NIF scaffold, MLX prebuilt fetch, tensor round-trip.

RELEASE.md

@@ -8,6 +8,57 @@
 
 ## Added
 
+- M22 — Compile-time debug flags. Two opt-in `Application.compile_env`
+  flags re-enable runtime assertions on hot paths that GPU backends
+  skip by default. Both default to `false` with zero runtime cost
+  when off — the gated branches are dead-code eliminated and no
+  `Emily.Backend.DebugHelpers` reference reaches the compiled BEAM.
+  - **`:debug_bounds_check`** — raises on out-of-range / negative
+    indices in `gather`, `take`, `take_along_axis`, `indexed_add`,
+    `indexed_put`. Catches the class of bug where a vocab-N
+    tokenizer paired with a smaller embedding table silently emits
+    garbage (sometimes NaN) through `mx::take` / `mx::gather`. One
+    gate each in `gather/4` (both branches), `take/3`,
+    `take_along_axis/3`, and `apply_scatter/8` — the shared
+    indexed_add/indexed_put helper.
+  - **`:debug_detect_nan_inf`** — scans results of `matmul`,
+    `fast_rms_norm`, `fast_layer_norm`, and both
+    `fast_scaled_dot_product_attention` variants for NaN / Inf.
+    Surfaces numerics failures at the producing op rather than
+    downstream as `loss = NaN`. Standalone softmax has no Backend
+    callback in Emily (Axon-composed); only fused SDPA softmax is
+    scanned. Documented alongside the flag.
+  - **Mechanism.** Module attributes bound via
+    `Application.compile_env/3` default to `false`; the `if @flag`
+    gates fold at compile time, and
+    `Emily.Backend.DebugHelpers.check_bounds!` /
+    `check_nan_inf!` are never referenced from
+    `Emily.Backend.beam`. A zero-cost verification test inspects
+    the disassembled BEAM (`:beam_disasm.file/1`) and asserts no
+    reference survives under the default-off build — guards
+    against a future refactor accidentally unconditionalising the
+    gate.
+  - **Each assertion is a GPU sync** (one reduction + scalar
+    readback per gated op). Breaks lazy-graph fusion; only enable
+    in dev / CI, never in release.
+  - **New files**: `lib/emily/backend/debug_helpers.ex` (assertion
+    bodies), `config/config.exs` + `config/test.exs` (first config
+    files in the project — previously no `config/` directory),
+    `test/support/debug_fixture.ex` (exercises the gate→helper
+    composition with its own fixture-only `compile_env` flags),
+    `test/emily/debug_flags_test.exs` (15 tests: direct helper,
+    production negative control, fixture composition, zero-cost
+    verification).
+  - **Emily's own `config/test.exs` keeps prod flags default-false**
+    so `mix test` doesn't pay the GPU-sync cost on every run and we
+    don't mask perf regressions. Fixture-only flags
+    (`test_fixture_debug_bounds_check`,
+    `test_fixture_debug_detect_nan_inf`) are set to `true` to drive
+    the composition tests.
+  - **Docs**: `Emily` moduledoc gains a "Debug assertions" section
+    with the worked opt-in snippet; README gets a matching brief
+    after `## Observability`.
+
 - M18 — Observability & fallback telemetry. Makes silent
   `via_binary` round-trips and long-running memory drift observable
   without changing any op semantics.

config/config.exs

@@ -0,0 +1,5 @@
+import Config
+
+if File.exists?(Path.join(__DIR__, "#{config_env()}.exs")) do
+  import_config "#{config_env()}.exs"
+end

config/test.exs

@@ -0,0 +1,11 @@
+import Config
+
+# Production M22 flags stay default-false even in Emily's own tests so
+# `mix test` doesn't pay GPU-sync cost on every run and we don't mask
+# perf regressions. Fixture-only flags (set to `true`) drive the
+# gate→helper composition tests in test/support/debug_fixture.ex.
+config :emily,
+  debug_bounds_check: false,
+  debug_detect_nan_inf: false,
+  test_fixture_debug_bounds_check: true,
+  test_fixture_debug_detect_nan_inf: true

lib/emily.ex

@@ -9,6 +9,43 @@ defmodule Emily do
   The M0 surface is intentionally tiny: build a tensor from a binary,
   round-trip it back, and inspect shape/dtype. This is the narrowest
   slice that proves the NIF + MLX linking is healthy.
+
+  ## Debug assertions
+
+  Emily supports two compile-time flags that enable runtime assertions
+  on hot paths. Both default to `false` with zero runtime cost — the
+  guarded branches are dead-code-eliminated by the Elixir compiler
+  when the flag is `false`.
+
+      # config/dev.exs  (opt in only in development / CI)
+      import Config
+      config :emily,
+        debug_bounds_check: true,
+        debug_detect_nan_inf: true
+
+    * `:debug_bounds_check` — assert indices are in range for
+      `gather`, `take`, `take_along_axis`, `indexed_add`, and
+      `indexed_put`. Raises `ArgumentError` on out-of-range or
+      negative indices. GPU backends (Emily, EXLA, Torch-CUDA,
+      JAX-GPU) don't bounds-check by default — an OOB index gets
+      whatever bytes happen to live at that memory address, which
+      can silently produce `NaN` scores that propagate through
+      softmax. Turning this on in CI catches the bug class at the
+      offending op.
+
+    * `:debug_detect_nan_inf` — scan results of `matmul`,
+      `fast_rms_norm`, `fast_layer_norm`, and the two
+      `fast_scaled_dot_product_attention` variants for NaN/Inf.
+      Raises `ArgumentError` on detection. Useful during training
+      so numerics failures surface at the op that produced them
+      rather than downstream as `loss = NaN`. Standalone softmax
+      has no backend callback in Emily (it's Axon-composed from
+      `exp` / `sum` / `divide`); only the fused SDPA softmax is
+      scanned.
+
+  Each assertion forces a small MLX reduction plus a scalar readback
+  on the worker — a sync point that breaks lazy-graph fusion and
+  adds noticeable overhead. Leave the flags off in release builds.
   """
 
   alias Emily.Native

lib/emily/backend.ex

@@ -31,9 +31,17 @@ defmodule Emily.Backend do
   defstruct [:ref]
 
   alias Emily.Backend, as: B
+  alias Emily.Backend.DebugHelpers
   alias Emily.Native
   alias Nx.Tensor, as: T
 
+  # Compile-time debug flags (M22). Default `false` so the `if @flag`
+  # gates below fold at compile time — the `DebugHelpers.*` references
+  # never appear in this module's BEAM when the flags are off. See
+  # `Emily`'s moduledoc and `test/emily/debug_flags_test.exs`.
+  @debug_bounds_check Application.compile_env(:emily, :debug_bounds_check, false)
+  @debug_detect_nan_inf Application.compile_env(:emily, :debug_detect_nan_inf, false)
+
   @typep tensor :: T.t()
   @typep ref :: reference()
 
@@ -665,11 +673,19 @@ defmodule Emily.Backend do
       match?([_], axes) ->
         [axis] = axes
         idx_ref = Native.astype(w, ref(indices), {:s, 32})
+
+        if @debug_bounds_check,
+          do: DebugHelpers.check_bounds!(:gather, input.shape, [idx_ref], [axis], w)
+
         r = Native.take(w, ref(input), idx_ref, axis)
         Native.reshape(w, r, Tuple.to_list(out.shape)) |> wrap(out, w)
 
       scatter_gather_compatible?(indices_shape, axes) ->
         idx_refs = split_indices_per_axis(ref(indices), indices_shape, length(axes), w)
+
+        if @debug_bounds_check,
+          do: DebugHelpers.check_bounds!(:gather, input.shape, idx_refs, axes, w)
+
         slice_sizes = slice_sizes_for_gather(input.shape, axes)
 
         r = Native.gather(w, ref(input), idx_refs, axes, slice_sizes)
@@ -833,6 +849,7 @@ defmodule Emily.Backend do
     rb = Native.reshape(w, rb, [b_prod, k_prod, n])
 
     r = Native.matmul(w, ra, rb)
+    if @debug_detect_nan_inf, do: DebugHelpers.check_nan_inf!(:matmul, r, w)
     Native.reshape(w, r, shape_list(out_shape)) |> wrap(out, w)
   end
 
@@ -932,6 +949,10 @@ defmodule Emily.Backend do
     w = worker()
     axis = opts[:axis] || 0
     idx_ref = Native.astype(w, ref(indices), {:s, 32})
+
+    if @debug_bounds_check,
+      do: DebugHelpers.check_bounds!(:take, input.shape, [idx_ref], [axis], w)
+
     Native.take(w, ref(input), idx_ref, axis) |> wrap(out, w)
   end
 
@@ -940,6 +961,10 @@ defmodule Emily.Backend do
     w = worker()
     axis = opts[:axis] || 0
     idx_ref = Native.astype(w, ref(indices), {:s, 32})
+
+    if @debug_bounds_check,
+      do: DebugHelpers.check_bounds!(:take_along_axis, input.shape, [idx_ref], [axis], w)
+
     Native.take_along_axis(w, ref(input), idx_ref, axis) |> wrap(out, w)
   end
 
@@ -1303,6 +1328,7 @@ defmodule Emily.Backend do
     if scatter_gather_compatible?(indices_shape, axes) do
       w = worker()
       idx_refs = split_indices_per_axis(ref(indices), indices_shape, length(axes), w)
+      if @debug_bounds_check, do: DebugHelpers.check_bounds!(op, t.shape, idx_refs, axes, w)
       updates_shape = updates_shape_for_scatter(indices_shape, t.shape, axes)
       updates_ref = Native.reshape(w, ref(updates), updates_shape)
 
@@ -1482,13 +1508,17 @@ defmodule Emily.Backend do
   @doc false
   def fast_rms_norm(%T{} = out, x, weight, opts) do
     w = worker()
-    Native.fast_rms_norm(w, ref(x), ref(weight), opts[:eps] * 1.0) |> wrap(out, w)
+    r = Native.fast_rms_norm(w, ref(x), ref(weight), opts[:eps] * 1.0)
+    if @debug_detect_nan_inf, do: DebugHelpers.check_nan_inf!(:fast_rms_norm, r, w)
+    wrap(r, out, w)
   end
 
   @doc false
   def fast_layer_norm(%T{} = out, x, weight, bias, opts) do
     w = worker()
-    Native.fast_layer_norm(w, ref(x), ref(weight), ref(bias), opts[:eps] * 1.0) |> wrap(out, w)
+    r = Native.fast_layer_norm(w, ref(x), ref(weight), ref(bias), opts[:eps] * 1.0)
+    if @debug_detect_nan_inf, do: DebugHelpers.check_nan_inf!(:fast_layer_norm, r, w)
+    wrap(r, out, w)
   end
 
   @doc false
@@ -1534,31 +1564,41 @@ defmodule Emily.Backend do
     w = worker()
     mask_mode = if opts[:causal], do: "causal", else: ""
 
-    Native.fast_scaled_dot_product_attention(
-      w,
-      ref(q),
-      ref(k),
-      ref(v),
-      opts[:scale] * 1.0,
-      mask_mode,
-      []
-    )
-    |> wrap(out, w)
+    r =
+      Native.fast_scaled_dot_product_attention(
+        w,
+        ref(q),
+        ref(k),
+        ref(v),
+        opts[:scale] * 1.0,
+        mask_mode,
+        []
+      )
+
+    if @debug_detect_nan_inf,
+      do: DebugHelpers.check_nan_inf!(:fast_scaled_dot_product_attention, r, w)
+
+    wrap(r, out, w)
   end
 
   @doc false
   def fast_scaled_dot_product_attention_with_mask(%T{} = out, q, k, v, mask, opts) do
     w = worker()
 
-    Native.fast_scaled_dot_product_attention(
-      w,
-      ref(q),
-      ref(k),
-      ref(v),
-      opts[:scale] * 1.0,
-      "array",
-      [ref(mask)]
-    )
-    |> wrap(out, w)
+    r =
+      Native.fast_scaled_dot_product_attention(
+        w,
+        ref(q),
+        ref(k),
+        ref(v),
+        opts[:scale] * 1.0,
+        "array",
+        [ref(mask)]
+      )
+
+    if @debug_detect_nan_inf,
+      do: DebugHelpers.check_nan_inf!(:fast_scaled_dot_product_attention_with_mask, r, w)
+
+    wrap(r, out, w)
   end
 end

lib/emily/backend/debug_helpers.ex

@@ -0,0 +1,50 @@
+defmodule Emily.Backend.DebugHelpers do
+  @moduledoc false
+  # Assertion bodies for the M22 compile-time debug flags. Always
+  # compiled; `Emily.Backend` only calls into here from `if @flag` gates
+  # that fold to nothing when the attribute is the literal `false`, so
+  # the reference to this module never lands in `Emily.Backend.beam`
+  # under the default-off build.
+
+  alias Emily.Native
+
+  @doc false
+  @spec check_bounds!(atom(), tuple(), [reference()], [integer()], reference()) :: :ok
+  def check_bounds!(op, input_shape, idx_refs, axes, w) do
+    axes
+    |> Enum.zip(idx_refs)
+    |> Enum.each(fn {axis, idx_ref} ->
+      dim = elem(input_shape, axis)
+      red_axes = Enum.to_list(0..(length(Native.shape(idx_ref)) - 1))
+
+      max_ref = Native.astype(w, Native.max(w, idx_ref, red_axes, false), {:s, 32})
+      min_ref = Native.astype(w, Native.min(w, idx_ref, red_axes, false), {:s, 32})
+      <<max_i::signed-integer-32-native>> = Native.to_binary(w, max_ref)
+      <<min_i::signed-integer-32-native>> = Native.to_binary(w, min_ref)
+
+      cond do
+        min_i < 0 ->
+          raise ArgumentError,
+                "#{op}: index #{min_i} on axis #{axis} is negative (dim=#{dim})"
+
+        max_i >= dim ->
+          raise ArgumentError,
+                "#{op}: index #{max_i} on axis #{axis} out of range (dim=#{dim})"
+
+        true ->
+          :ok
+      end
+    end)
+  end
+
+  @doc false
+  @spec check_nan_inf!(atom(), reference(), reference()) :: :ok
+  def check_nan_inf!(op, result_ref, w) do
+    bad = Native.logical_or(w, Native.isnan(w, result_ref), Native.isinf(w, result_ref))
+    axes = Enum.to_list(0..(length(Native.shape(bad)) - 1))
+    scalar = Native.any(w, bad, axes, false)
+    <<flag::unsigned-integer-8>> = Native.to_binary(w, scalar)
+    if flag == 1, do: raise(ArgumentError, "#{op}: produced NaN or Inf")
+    :ok
+  end
+end