From bba71c1f02e7b78033771d3f0defbf3d907b9c49 Mon Sep 17 00:00:00 2001 From: edenamram <74473448+edenamram@users.noreply.github.com> Date: Sun, 14 Jun 2026 15:05:33 +0300 Subject: [PATCH 1/4] Fix int32 overflow in elem_count for shapes with >INT32_MAX elements --- emlx/c_src/emlx_nif.cpp | 2 +- emlx/test/emlx_test.exs | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/emlx/c_src/emlx_nif.cpp b/emlx/c_src/emlx_nif.cpp index b62948d..f18a408 100644 --- a/emlx/c_src/emlx_nif.cpp +++ b/emlx/c_src/emlx_nif.cpp @@ -218,7 +218,7 @@ NIF(to_blob) { } uint64_t elem_count(std::vector shape) { - return std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<>{}); + return std::accumulate(shape.begin(), shape.end(), uint64_t{1}, std::multiplies{}); } NIF(from_blob) { diff --git a/emlx/test/emlx_test.exs b/emlx/test/emlx_test.exs index 53cc1ae..d42e14c 100644 --- a/emlx/test/emlx_test.exs +++ b/emlx/test/emlx_test.exs @@ -170,4 +170,23 @@ defmodule EMLXTest do end end end + + describe "large tensors (element count > INT32_MAX)" do + # Regression: elem_count overflowed int32 for shapes whose element count + # exceeds INT32_MAX, causing "Binary size is too small" on valid binaries. + @tag :large_tensor + test "from_binary accepts shape whose element count exceeds INT32_MAX" do + # Reshape on BinaryBackend first — Nx.from_binary creates a flat 1D tensor + # whose single dimension would also exceed INT32_MAX. + binary = :binary.copy(<<7>>, 2_147_483_648) + + t = + Nx.from_binary(binary, :u8, backend: Nx.BinaryBackend) + |> Nx.reshape({2, 1_073_741_824}) + |> Nx.backend_transfer(EMLX.Backend) + + assert Nx.shape(t) == {2, 1_073_741_824} + assert Nx.to_number(t[0][0]) == 7 + end + end end From ec53ef30a81e1333a677bd8826248695b186803d Mon Sep 17 00:00:00 2001 From: Paulo Valente <16843419+polvalente@users.noreply.github.com> Date: Sun, 14 Jun 2026 17:11:42 -0300 Subject: [PATCH 2/4] simplify test --- emlx/test/emlx_test.exs | 31 ++++++++++++++----------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/emlx/test/emlx_test.exs b/emlx/test/emlx_test.exs index d42e14c..5784787 100644 --- a/emlx/test/emlx_test.exs +++ b/emlx/test/emlx_test.exs @@ -171,22 +171,19 @@ defmodule EMLXTest do end end - describe "large tensors (element count > INT32_MAX)" do - # Regression: elem_count overflowed int32 for shapes whose element count - # exceeds INT32_MAX, causing "Binary size is too small" on valid binaries. - @tag :large_tensor - test "from_binary accepts shape whose element count exceeds INT32_MAX" do - # Reshape on BinaryBackend first — Nx.from_binary creates a flat 1D tensor - # whose single dimension would also exceed INT32_MAX. - binary = :binary.copy(<<7>>, 2_147_483_648) - - t = - Nx.from_binary(binary, :u8, backend: Nx.BinaryBackend) - |> Nx.reshape({2, 1_073_741_824}) - |> Nx.backend_transfer(EMLX.Backend) - - assert Nx.shape(t) == {2, 1_073_741_824} - assert Nx.to_number(t[0][0]) == 7 - end + # Regression: elem_count overflowed int32 for shapes whose element count + # exceeds INT32_MAX, causing "Binary size is too small" on valid binaries. + test "from_binary accepts shape whose element count exceeds INT32_MAX" do + # Reshape on BinaryBackend first — Nx.from_binary creates a flat 1D tensor + # whose single dimension would also exceed INT32_MAX. + int_32_max = 2 ** 31 + binary = String.duplicate(<<7>>, int_32_max) + shape = {2, div(int_32_max, 2)} + + out = Nx.template(shape, :u8) + t = EMLX.Backend.from_binary(out, binary, backend: EMLX.Backend) + + assert Nx.shape(t) == shape + assert_equal(Nx.all(Nx.equal(t, 7)), 1) end end From decb36d1b2922a8ef59a3d7dca8acf22705b71d5 Mon Sep 17 00:00:00 2001 From: Paulo Valente <16843419+polvalente@users.noreply.github.com> Date: Sun, 14 Jun 2026 17:23:41 -0300 Subject: [PATCH 3/4] fix: ci --- .github/workflows/emlx.yml | 2 +- emlx/test/emlx_test.exs | 1 + emlx/test/test_helper.exs | 4 ++-- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/emlx.yml b/.github/workflows/emlx.yml index 6e50dad..8fd1392 100644 --- a/.github/workflows/emlx.yml +++ b/.github/workflows/emlx.yml @@ -72,7 +72,7 @@ jobs: mix test --warnings-as-errors macos: - name: macOS (${{ matrix.job.elixir }}, ${{ matrix.job.otp }}) + name: macOS ${{ matrix.job.gpu && 'gpu' || 'cpu' }} (${{ matrix.job.elixir }}, ${{ matrix.job.otp }}) runs-on: macos-26 strategy: fail-fast: false diff --git a/emlx/test/emlx_test.exs b/emlx/test/emlx_test.exs index 5784787..8d2a33a 100644 --- a/emlx/test/emlx_test.exs +++ b/emlx/test/emlx_test.exs @@ -173,6 +173,7 @@ defmodule EMLXTest do # Regression: elem_count overflowed int32 for shapes whose element count # exceeds INT32_MAX, causing "Binary size is too small" on valid binaries. + @tag :large_memory test "from_binary accepts shape whose element count exceeds INT32_MAX" do # Reshape on BinaryBackend first — Nx.from_binary creates a flat 1D tensor # whose single dimension would also exceed INT32_MAX. diff --git a/emlx/test/test_helper.exs b/emlx/test/test_helper.exs index 4cf03a2..762dfa3 100644 --- a/emlx/test/test_helper.exs +++ b/emlx/test/test_helper.exs @@ -51,8 +51,8 @@ distributed_exclude = gpu_exclude = case EMLX.NIF.command_queue_new(:gpu) do - {:ok, _} -> [] + {:ok, _} -> [:large_memory] {:error, _} -> [:metal] end -ExUnit.start(exclude: distributed_exclude ++ gpu_exclude) +ExUnit.start(exclude: distributed_exclude ++ gpu_exclude ++ large_memory_exclude) From f7553491cecc8b8551a02a9b6d62b79859506409 Mon Sep 17 00:00:00 2001 From: Paulo Valente <16843419+polvalente@users.noreply.github.com> Date: Sun, 14 Jun 2026 17:38:39 -0300 Subject: [PATCH 4/4] fix --- emlx/test/test_helper.exs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/emlx/test/test_helper.exs b/emlx/test/test_helper.exs index 762dfa3..a6ed859 100644 --- a/emlx/test/test_helper.exs +++ b/emlx/test/test_helper.exs @@ -55,4 +55,4 @@ gpu_exclude = {:error, _} -> [:metal] end -ExUnit.start(exclude: distributed_exclude ++ gpu_exclude ++ large_memory_exclude) +ExUnit.start(exclude: distributed_exclude ++ gpu_exclude)