Give a zero remainder the sign of the divisor

mlx/backend/cpu/simd/accelerate_simd.h

@@ -242,8 +242,17 @@ Simd<T, N> remainder(Simd<T, N> a, Simd<T, N> b) {
     r = a - b * (a / b);
   }
   if constexpr (is_signed_v<T>) {
-    auto mask = r != 0 && (r < 0 != b < 0);
+    auto nonzero = r != 0;
+    auto mask = nonzero && (r < 0 != b < 0);
     r = select(mask, r + b, r);
+    if constexpr (!std::is_integral_v<T>) {
+      // A zero remainder takes the sign of the divisor, which `r + b` cannot
+      // give it: adding to a zero would return `b` itself. See base_simd.h.
+      // Selecting on `b < 0` rather than multiplying by zero keeps an infinite
+      // divisor working, since `inf * 0` is NaN but `remainder(0, inf)` is 0.
+      Simd<T, N> zero = static_cast<T>(0);
+      r = select(nonzero, r, select(b < 0, -zero, zero));
+    }
   }
   return r;
 }

mlx/backend/cpu/simd/base_simd.h

@@ -214,8 +214,15 @@ Simd<T, 1> remainder(Simd<T, 1> a_, Simd<T, 1> b_) {
     r = std::remainder(a, b);
   }
   if constexpr (is_signed_v<T>) {
-    if (r != 0 && (r < 0 != b < 0)) {
-      r += b;
+    if (r != 0) {
+      if (r < 0 != b < 0) {
+        r += b;
+      }
+    } else if constexpr (is_floating_point_v<T>) {
+      // A zero remainder takes the sign of the divisor. `r += b` cannot do it,
+      // since adding to a zero would return `b` itself, so set the sign
+      // directly. Integers have no signed zero and need nothing here.
+      r = static_cast<T>(std::copysign(0.0f, static_cast<float>(b)));
     }
   }
   return r;

mlx/backend/cuda/device/binary_ops.cuh

@@ -48,8 +48,16 @@ struct Remainder {
       return x % y;
     } else {
       T r = cuda::std::fmod(x, y);
-      if (r != 0 && (r < 0 != y < 0)) {
-        r = r + y;
+      if (r != 0) {
+        if (r < 0 != y < 0) {
+          r = r + y;
+        }
+      } else {
+        // A zero remainder takes the sign of the divisor. `r + y` cannot give
+        // it that sign, since adding to a zero returns `y` itself. Written as a
+        // comparison rather than copysign because the half types do not have a
+        // copysign overload. A divisor of zero cannot reach here: fmod is NaN.
+        r = y < T(0) ? T(-0.0f) : T(0.0f);
       }
       return r;
     }

mlx/backend/metal/kernels/binary_ops.h

@@ -58,8 +58,14 @@ struct Remainder {
   template <typename T>
   metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T x, T y) thread {
     T r = fmod(x, y);
-    if (r != 0 && (r < 0 != y < 0)) {
-      r += y;
+    if (r != 0) {
+      if (r < 0 != y < 0) {
+        r += y;
+      }
+    } else {
+      // A zero remainder takes the sign of the divisor. `r += y` cannot give it
+      // that sign, since adding to a zero returns `y` itself.
+      r = metal::copysign(T(0), y);
     }
     return r;
   }

python/tests/test_ops.py

@@ -401,6 +401,50 @@ def test_remainder(self):
         z = -mx.ones(64) % mx.full(64, 2)
         self.assertTrue(mx.array_equal(z, mx.ones(64)))
 
+    def test_remainder_signed_zero(self):
+        # A zero remainder takes the sign of the divisor. `==` cannot check this,
+        # since 0.0 == -0.0, so every assertion here goes through signbit.
+        for dt in [mx.float32, mx.float16, mx.bfloat16]:
+            # Length 9 puts the first 8 elements in the vectorized body and the
+            # last one in the scalar residual. All nine inputs are the same, so
+            # all nine signs have to be.
+            got = np.array(
+                mx.remainder(
+                    mx.full((9,), -0.0, dtype=dt), mx.full((9,), 3.0, dtype=dt)
+                ).astype(mx.float32)
+            )
+            self.assertEqual(np.signbit(got).tolist(), [False] * 9)
+
+            # A strided view has to agree with a contiguous array of the same
+            # values, since the layout carries no numerical information.
+            a = mx.full((16,), -0.0, dtype=dt)
+            b = mx.full((16,), 3.0, dtype=dt)
+            strided = np.array(mx.remainder(a[::2], b[::2]).astype(mx.float32))
+            self.assertEqual(np.signbit(strided).tolist(), [False] * 8)
+
+            av = [5.0, -5.0, 6.0, -6.0, 0.0, -0.0, 7.0, -7.0]
+            bv = [-5.0, 5.0, -3.0, 3.0, -1.0, 3.0, -7.0, 7.0]
+            got = np.array(
+                mx.remainder(mx.array(av, dtype=dt), mx.array(bv, dtype=dt)).astype(
+                    mx.float32
+                )
+            )
+            expected = np.remainder(
+                np.array(av, dtype=np.float32), np.array(bv, dtype=np.float32)
+            )
+            self.assertTrue(np.array_equal(got, expected))
+            self.assertEqual(np.signbit(got).tolist(), np.signbit(expected).tolist())
+
+            # The operator goes through the same kernel. mx.divmod does not; it
+            # truncates, which is #4119.
+            r = mx.array([5.0], dtype=dt) % mx.array([-5.0], dtype=dt)
+            self.assertTrue(np.signbit(np.array(r.astype(mx.float32))).item())
+
+        # Integers have no signed zero, so the correction must not touch them.
+        for dt in [mx.int32, mx.int64]:
+            z = mx.remainder(mx.array([5, -5, 6], dtype=dt), mx.array([-5, 5, -3], dt))
+            self.assertEqual(z.tolist(), [0, 0, 0])
+
     def test_comparisons(self):
         a = mx.array([0.0, 1.0, 5.0])
         b = mx.array([-1.0, 2.0, 5.0])