Fix BinaryBackend.select_and_scatter crash with :same padding (#1675)Fix/select and scatter negative padding#1676
Conversation
| absolute_index = | ||
| padded_absolute_index | ||
| |> Enum.zip(low_pads) | ||
| |> Enum.map(fn {idx, lo} -> idx - lo end) |
There was a problem hiding this comment.
| absolute_index = | |
| padded_absolute_index | |
| |> Enum.zip(low_pads) | |
| |> Enum.map(fn {idx, lo} -> idx - lo end) | |
| absolute_index = Enum.zip_with(padded_absolute_index, low_pads, &-/2) |
| Enum.zip(index, original_dims) | ||
| |> Enum.all?(fn {idx, dim} -> idx >= 0 and idx < dim end) |
There was a problem hiding this comment.
| Enum.zip(index, original_dims) | |
| |> Enum.all?(fn {idx, dim} -> idx >= 0 and idx < dim end) | |
| Enum.zip_reduce(index, original_dims, true, fn idx, dim, acc -> acc and idx >= 0 and idx < dim end) |
|
|
||
| # With kernel=3, stride=2, same padding on 4x4: | ||
| # output is 2x2, max positions are at bottom-right of each window | ||
| # Gradient should be 1.0 at the max positions, 0.0 elsewhere | ||
| assert Nx.shape(lhs) == {1, 4, 4, 1} | ||
| assert Nx.type(lhs) == {:f, 32} | ||
|
|
||
| # All gradient values should be finite (no NaN/Inf) | ||
| assert Nx.all(Nx.is_nan(lhs) |> Nx.logical_not()) == Nx.tensor(1, type: {:u, 8}) | ||
|
|
||
| # Gradient should sum to the number of output elements (each max gets grad 1.0) | ||
| assert_all_close(Nx.sum(lhs), Nx.tensor(4.0)) |
There was a problem hiding this comment.
I think the tensor is small enough that there should be an asser_equal or assert_all_close assertion on the lhs values directly.
Also, there should be new tests for scatter_max directly besides this one.
There was a problem hiding this comment.
Also, for this assertion, we should double check that jax returns the same values for both the defined function as well as a compound function (say we put the input through Nx.cos and the output through Nx.sin to ensure non-linearity of the grad propagation -- this should be a new test)
polvalente
left a comment
polvalente
left a comment
There was a problem hiding this comment.
Thanks a lot for taking this on. Code generally looks good, I just think we should have more tests :)
|
no problem! ran into it building another repo, also filed #1679 cause that popped up while trying to add more tests |
…#1675) Apply reviewer code style suggestions (Enum.zip_with, Enum.zip_reduce) and expand test coverage: exact gradient assertions, window_min/max/sum with same/explicit/asymmetric padding, 1D/2D/4D shapes, kernel>input, kernel==input, stride>kernel, and non-linear composition (cos→max→sin). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
blasphemetheus
force-pushed
the
fix/select-and-scatter-negative-padding
branch
from
122cf40 to
43d377f
Compare
|
Thanks a lot! I'm looking through all of your issues/PRs :) |
2 participants
Summary
select_and_scatter/8computesabsolute_indexin the padded coordinate space but uses it against the unpaddedoutput_shape, producing negative offsets that crashList.duplicate/2.Fix: subtract low-padding from indices to map back to original coordinates, and filter out indices that land in the padding region.
Changes
nx/lib/nx/binary_backend.exabsolute_indexafter computing it in padded spaceoutput_windowsto discard indices outside original tensor boundsnx/test/nx/defn/grad_test.exsgrad(sum(window_max(x, {1,3,3,1}, padding: :same, strides: [1,2,2,1])))on a{1,4,4,1}input — verifies output shape, finiteness, and gradient sum equals number of output elementsTest Results
Closes #1675