test: add regression test for large tensors exceeding INT32_MAX

Author: edenamram

emlx/test/emlx_test.exs

@@ -170,4 +170,33 @@ defmodule EMLXTest do
       end
     end
   end
+
+  describe "large tensors (element count > INT32_MAX)" do
+    # Regression: elem_count used std::accumulate with an `int` initial value,
+    # so products > INT32_MAX overflowed to a huge uint64 and the binary-size
+    # guard rejected valid binaries with "Binary size is too small".
+    # Real-world trigger: Gemma-4's embed_tokens_per_layer {262144, 8960}
+    #   = 2_348_810_240 elements > INT32_MAX (2_147_483_647).
+    @tag :large_tensor
+    test "from_binary accepts shape whose element count exceeds INT32_MAX" do
+      # 2 × 1_073_741_824 = 2_147_483_648 > INT32_MAX.
+      # :u8 (1 byte/element) keeps total allocation at ~2 GB.
+      #
+      # We reshape on BinaryBackend first (pure Elixir, no NIF) because
+      # Nx.from_binary always creates a 1D tensor — and a single dimension of
+      # 2_147_483_648 also exceeds INT32_MAX, which breaks SHAPE_PARAM before
+      # even reaching elem_count.  The 2D shape {2, 1_073_741_824} has both
+      # dims within int32 range; only their product overflows.
+      n = 2_147_483_648
+      binary = :binary.copy(<<7>>, n)
+
+      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