Nx.runtime_call receives stride-0 tensors with only the first element intact under compiler: EMLX

[Codcore]

What happens

A tensor with zero strides — anything produced by Nx.broadcast/2 — arrives at an
Nx.runtime_call/4 callback with only its first element intact when the enclosing
function is compiled with compiler: EMLX. The remaining elements read as zeros.

The same call evaluated eagerly is correct.

stride-0 inputs (Nx.broadcast):
  broadcast u32                      MISMATCH
      eager:    [3, 3, 3, 3]
      compiled: [3, 0, 0, 0]
  broadcast f32                      MISMATCH
      eager:    [2.5, 2.5, 2.5, 2.5]
      compiled: [2.5, 0.0, 0.0, 0.0]

Where the corruption happens

On the way in, not on the way out. A callback that sums its argument shows it
directly:

sum of [2.5, 2.5, 2.5, 2.5], expected 10.0
  eager:    10.0
  compiled: 2.5

Printing from inside the callback confirms the kernel itself receives
[2.5, 0.0, 0.0, 0.0].

What is not affected

• Pure Nx operations on the same tensors (Nx.add, Nx.sum, Nx.take) are correct
  under the same compiler.
• EMLX's own fused kernels are correct: EMLX.Fast.rms_norm/3 and
  EMLX.Fast.swiglu/2 handle broadcast arguments fine. They route through
  emlx_metadata/4, which is not part of the public surface, so a third-party
  fused op has no way to take the same path.

Nx.backend_copy/2 does not help — the copy keeps the same layout.

Why it matters

Nx.runtime_call/4 is the extension point for adding a fused kernel from outside
EMLX. Broadcast tensors are ordinary values in model code (a scalar gate expanded
over a batch, an all-ones norm weight, a repeated index vector), so the failure is
silent and data-dependent: the shapes are right, no error is raised, and the numbers
are merely wrong.

Reproduction

Self-contained script, no dependencies beyond {:emlx, "~> 0.1"}:

mix run runtime_call_broadcast_repro.exs

Every case should print ok.

Environment

• emlx 0.4.1 (Hex release, unmodified)
• nx 0.13
• libmlx 0.31.2, arm64-apple-darwin
• Elixir 1.20.3, Erlang/OTP 29
• macOS 26.5, Apple Silicon

Possible fix

Materialise non-contiguous inputs before handing them to the callback in the
compiled path, the way the eager path already does. Alternatively, document that
Nx.runtime_call arguments must be contiguous and expose whatever emlx_metadata
does for the built-in kernels.

runtime_call_broadcast_repro.exs

# Minimal reproduction: Nx.runtime_call receives corrupted stride-0 tensors
# under compiler: EMLX.
#
#   mix run runtime_call_broadcast_repro.exs
#
# Expected: every case prints "ok".
# Actual:   broadcast inputs reach the callback with only the first element
#           intact; the rest read as zeros.

Nx.default_backend({EMLX.Backend, device: :gpu})

defmodule Repro do
  import Nx.Defn

  # The callback prints exactly what the kernel receives.
  def report({t}, _opts) do
    IO.puts("      kernel received: #{inspect(Nx.to_flat_list(t))}")
    t
  end

  deftransform identity(t) do
    if match?(%Nx.Tensor{data: %Nx.Defn.Expr{}}, t) do
      Nx.runtime_call(Nx.to_template(t), {t}, [], &report/2)
    else
      report({t}, [])
    end
  end

  def sum_cb({t}, _opts), do: Nx.sum(t)

  deftransform total(t) do
    if match?(%Nx.Tensor{data: %Nx.Defn.Expr{}}, t) do
      Nx.runtime_call(Nx.template({}, Nx.type(t)), {t}, [], &sum_cb/2)
    else
      sum_cb({t}, [])
    end
  end
end

gpu = fn t -> Nx.backend_transfer(t, {EMLX.Backend, device: :gpu}) end

check = fn label, tensor ->
  eager = Repro.identity(tensor) |> Nx.to_flat_list()
  compiled = Nx.Defn.jit(&Repro.identity/1, compiler: EMLX).(tensor) |> Nx.to_flat_list()

  IO.puts("  #{String.pad_trailing(label, 34)} #{if eager == compiled, do: "ok", else: "MISMATCH"}")

  if eager != compiled do
    IO.puts("      eager:    #{inspect(eager)}")
    IO.puts("      compiled: #{inspect(compiled)}")
  end
end

IO.puts("\ncontiguous inputs:")
check.("dense u32", gpu.(Nx.tensor([3, 3, 3, 3], type: :u32)))
check.("dense f32", gpu.(Nx.tensor([2.5, 2.5, 2.5, 2.5])))

IO.puts("\nstride-0 inputs (Nx.broadcast):")
check.("broadcast u32", gpu.(Nx.broadcast(Nx.u32(3), {4})))
check.("broadcast f32", gpu.(Nx.broadcast(2.5, {4})))
check.("broadcast f32, backend_copy", gpu.(Nx.broadcast(2.5, {4})) |> Nx.backend_copy({EMLX.Backend, device: :gpu}))

IO.puts("\nthe corruption is on the way in, not on the way out:")
b = gpu.(Nx.broadcast(2.5, {4}))
IO.puts("  sum of [2.5, 2.5, 2.5, 2.5], expected 10.0")
IO.puts("    eager:    #{Repro.total(b) |> Nx.to_number()}")
IO.puts("    compiled: #{Nx.Defn.jit(&Repro.total/1, compiler: EMLX).(b) |> Nx.to_number()}")

IO.puts("\npure Nx ops on the same tensors are unaffected:")
pure = fn label, f, args ->
  e = apply(f, args) |> Nx.to_flat_list()
  c = apply(Nx.Defn.jit(f, compiler: EMLX), args) |> Nx.to_flat_list()
  IO.puts("  #{String.pad_trailing(label, 34)} #{if e == c, do: "ok", else: "MISMATCH"}")
end

pure.("Nx.add with broadcast", &Nx.add/2, [gpu.(Nx.tensor([1.0, 2.0, 3.0])), gpu.(Nx.broadcast(1.0, {3}))])
pure.("Nx.sum of broadcast", &Nx.sum/1, [gpu.(Nx.broadcast(2.0, {5}))])
pure.("Nx.take with broadcast index", &Nx.take(&1, &2, axis: 0),
      [gpu.(Nx.iota({8, 2}, type: :f32)), gpu.(Nx.broadcast(Nx.u32(3), {4}))])

[polvalente]

A PR would be welcome! Small correction: if you look at EMLXAxon there's a pattern that EMLX exposes for plugins, which is the recommended way for adding fused acceleration.

The issue (likely an LLM) mentions runtime_call as the recommended path