Fix integer pow zeroing a whole SIMD vector on a negative exponent

[ayaangazali]

What is wrong

On the CPU Accelerate backend, mx.power on integers zeroes an entire SIMD vector when any single element has a negative exponent, so correct positive-exponent elements come back as 0.

import mlx.core as mx
with mx.stream(mx.cpu):
    base = mx.array([2] * 16, mx.int32)
    exp  = mx.array([3, 3, 3, -1] + [3] * 12, mx.int32)
    print(mx.power(base, exp).tolist())
# [0, 0, 0, 0, 0, 0, 0, 0, 8, 8, 8, 8, 8, 8, 8, 8]
# expected:
# [8, 8, 8, 0, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8]

Only the element with exponent -1 should be 0 (mlx defines integer-to-negative-power as 0). Instead the whole 8-lane int32 register that element lands in is zeroed, so 2 ** 3 silently returns 0 for its neighbors.

Why

simd::pow for integers bailed out for the whole vector as soon as one lane was negative:

Simd<T, N> res = 1;
// Raising an integer to a negative power is undefined
if (any(exp < 0)) {
  return 0;
}

any(exp < 0) is true if a single lane is negative, and the return 0 then applies to all N lanes. The scalar path (base_simd.h, N == 1) and the Metal/CUDA kernels are per element and are not affected, so this is specific to the vectorized Accelerate path and only shows up when negative and non-negative exponents share a register.

What this changes

Run the existing exponentiation loop unconditionally and mask to 0 only the lanes whose exponent is negative:

Simd<T, N> res = 1;
while (any(exp > 0)) {
  res = select((exp & 1) != 0, res * base, res);
  base = select(exp > 0, base * base, base);
  exp = exp >> 1;
}
if constexpr (std::is_signed_v<T>) {
  res = select(exp < 0, Simd<T, N>(0), res);
}

Negative lanes never satisfy exp > 0, so they stay at the identity through the loop and are set to 0 at the end; the arithmetic right shift keeps them negative so the final mask still selects them. The all-negative fast path is preserved (the while never runs), and the exp < 0 compare is guarded by if constexpr (std::is_signed_v<T>) so the unsigned instantiations do not trip -Wtype-limits under -DCMAKE_COMPILE_WARNING_AS_ERROR=ON.

Added a mixed-sign case to test_integer_power; it fails on main ([0, 0, 0, 0, 0, 0, 0, 0, 8, ...]) and passes with this change. Full C++ suite (251 cases) and test_ops pass.

---

Sent by a beginner running Claude Code, so I checked it carefully: reverted the header on a fresh build to confirm the new test fails without the fix, and reran the C++ and ops suites after restoring it. Please tell me if anything needs another pass.