From e702454aa567a9460513d2c5a11b6c6227382ece Mon Sep 17 00:00:00 2001 From: Bradley Lewis Fargo Date: Fri, 20 Mar 2026 01:34:21 -0500 Subject: [PATCH 1/2] Fix window_scatter_max/min crash on f64 tensors The scatter result was not cast to the output type before to_binary, causing a binary size mismatch for f64 (8-byte) tensors. Co-Authored-By: Claude Opus 4.6 (1M context) --- nx/lib/nx/binary_backend.ex | 2 +- nx/test/nx_test.exs | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/nx/lib/nx/binary_backend.ex b/nx/lib/nx/binary_backend.ex index ae9c3f6267..85a1eb9f7b 100644 --- a/nx/lib/nx/binary_backend.ex +++ b/nx/lib/nx/binary_backend.ex @@ -1675,7 +1675,7 @@ defmodule Nx.BinaryBackend do {acc_offset, acc_binary} -> num_vals_before = div(offset - acc_offset, output_size) vals_before = List.duplicate(init_binary, num_vals_before) - source_val = to_binary(value) + source_val = value |> Nx.as_type(output_type) |> to_binary() new_binary = :erlang.list_to_bitstring([vals_before, source_val]) {offset + output_size, <>} diff --git a/nx/test/nx_test.exs b/nx/test/nx_test.exs index 5d42216132..73d1b36568 100644 --- a/nx/test/nx_test.exs +++ b/nx/test/nx_test.exs @@ -1059,6 +1059,24 @@ defmodule NxTest do ]) end + test "computes window scatter max with f64" do + t = Nx.iota({6}, type: :f64) + s = Nx.iota({3}, type: :f64) + init = Nx.tensor(0.0, type: :f64) + result = Nx.window_scatter_max(t, s, init, {2}, strides: [2], padding: :valid) + assert Nx.type(result) == {:f, 64} + assert Nx.shape(result) == {6} + end + + test "computes window scatter min with f64" do + t = Nx.iota({6}, type: :f64) + s = Nx.iota({3}, type: :f64) + init = Nx.tensor(0.0, type: :f64) + result = Nx.window_scatter_min(t, s, init, {2}, strides: [2], padding: :valid) + assert Nx.type(result) == {:f, 64} + assert Nx.shape(result) == {6} + end + test "computes window reduce (sum of squares) with same padding" do t = Nx.iota({4, 4}, type: {:f, 32}) From d5e5d86b093a436e5cf6fd998c84c3998cb33447 Mon Sep 17 00:00:00 2001 From: Paulo Valente <16843419+polvalente@users.noreply.github.com> Date: Mon, 23 Mar 2026 23:33:33 -0300 Subject: [PATCH 2/2] refactor: simplify fix --- nx/lib/nx/binary_backend.ex | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/nx/lib/nx/binary_backend.ex b/nx/lib/nx/binary_backend.ex index 85a1eb9f7b..33ad544c4d 100644 --- a/nx/lib/nx/binary_backend.ex +++ b/nx/lib/nx/binary_backend.ex @@ -1631,11 +1631,7 @@ defmodule Nx.BinaryBackend do # Compute absolute index in padded space, then adjust back to # original (unpadded) coordinates by subtracting low-padding - padded_absolute_index = - anchor - |> Enum.zip(offset_from_anchor) - |> Enum.map(fn {x, y} -> x + y end) - + padded_absolute_index = Enum.zip_with(anchor, offset_from_anchor, &+/2) absolute_index = Enum.zip_with(padded_absolute_index, low_pads, &-/2) source_consumed = i * source_size @@ -1664,7 +1660,8 @@ defmodule Nx.BinaryBackend do |> Enum.group_by(&elem(&1, 1), &elem(&1, 0)) |> Enum.map(fn {index, value} -> offset = weighted_offset(output_weighted_shape, index) - {offset, Enum.reduce(value, init_value, scatter_fn)} + tensor = Enum.reduce(value, init_value, scatter_fn) + {offset, scalar_to_number(tensor)} end) |> Enum.sort_by(&elem(&1, 0)) @@ -1675,7 +1672,12 @@ defmodule Nx.BinaryBackend do {acc_offset, acc_binary} -> num_vals_before = div(offset - acc_offset, output_size) vals_before = List.duplicate(init_binary, num_vals_before) - source_val = value |> Nx.as_type(output_type) |> to_binary() + + source_val = + match_types [output_type] do + <> + end + new_binary = :erlang.list_to_bitstring([vals_before, source_val]) {offset + output_size, <>}