Fix int32 overflow in elem_count for tensors with >2.1B elements

[edenamram]

Problem

elem_count uses std::accumulate with an int initial value:

uint64_t elem_count(std::vector<int> shape) {
  return std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<>{});
}

When the initial accumulator value is int, all intermediate multiplications are performed in 32-bit signed arithmetic. For shapes whose element-count product exceeds INT32_MAX (~2.1B), the product overflows to a negative int, which is then cast to a large uint64_t value. The binary-size guard in from_blob then always fires:

** (RuntimeError) Binary size is too small for the requested shape

Reproduction

google/gemma-4-2b-it (and larger Gemma-4 variants) include a weight called embed_tokens_per_layer with shape {262144, 8960}:

262144 × 8960 = 2,348,810,240  >  INT32_MAX (2,147,483,647)

This causes elem_count({262144, 8960}) to return 18,446,744,071,762,394,560 instead of 2,348,810,240, making it impossible to load the model via EMLX.Backend.

Fix

Seed the accumulator with uint64_t{1} and use std::multiplies<uint64_t> so every partial product stays in 64-bit arithmetic:

uint64_t elem_count(std::vector<int> shape) {
  return std::accumulate(shape.begin(), shape.end(), uint64_t{1}, std::multiplies<uint64_t>{});
}

Context

Discovered while implementing Gemma-4 support in Bumblebee: elixir-nx/bumblebee#460.

🤖 Generated with Claude Code

[polvalente]

Thanks for the quick fix!