[dkuku]

Found using antigravity.

Problem

In Nx.BinaryBackend, aggregate_axes/3 is designed to expand read_size when reducing along trailing contiguous dimensions so contiguous chunks can be read in one operation.

However, aggregate_read/4 had a variable collision in its pattern match:

defp aggregate_read([{axis, weight} | shape], i, [i | axis], _size),
  do: aggregate_read(shape, i - 1, axis, axis * weight)

Because axis was bound in the first argument to an integer (the dimension size), [i | axis] required the tail of the axes list to match that same integer. Since the incoming axes argument is a normal list of integers ([i | axes]), this clause never matched and always fell through to the base clause:

defp aggregate_read(shape, _i, _axes, size), do: {shape, size}

Benchmark

For a 500×500 f32 matrix:

• sum along axis 1: 17.2 ms → 3.05 ms (~5.6x faster)
• Memory allocated: 170 KB → 51 KB (~70% reduction)

[josevalim]

💚 💙 💜 💛 ❤️