From 42574939b66202fabb1998619940bfef506ceeb1 Mon Sep 17 00:00:00 2001
From: Moinuddin Shaik <moinuddinmoin1357@gmail.com>
Date: Tue, 22 Sep 2026 23:40:18 +0530
Subject: [PATCH 1/2] Instantiate init reduce for unsigned sum/prod

---
 mlx/backend/metal/kernels/reduce.metal |  2 ++
 python/tests/test_reduce.py            | 30 ++++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/mlx/backend/metal/kernels/reduce.metal b/mlx/backend/metal/kernels/reduce.metal
index 3f5f91be49..d0f528f694 100644
--- a/mlx/backend/metal/kernels/reduce.metal
+++ b/mlx/backend/metal/kernels/reduce.metal
@@ -19,6 +19,8 @@ instantiate_init_reduce(or, bool_, bool, Or)
 #define instantiate_init_sum_prod(name, op)                 \
   instantiate_init_reduce(name, int32, int32_t, op)         \
   instantiate_init_reduce(name, int64, int64_t, op)         \
+  instantiate_init_reduce(name, uint32, uint32_t, op)       \
+  instantiate_init_reduce(name, uint64, uint64_t, op)       \
   instantiate_init_reduce(name, float16, float16_t, op)     \
   instantiate_init_reduce(name, bfloat16, bfloat16_t, op)   \
   instantiate_init_reduce(name, float32, float, op)         \
diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py
index 0b04038176..338d7cb7b0 100644
--- a/python/tests/test_reduce.py
+++ b/python/tests/test_reduce.py
@@ -187,6 +187,36 @@ def test_zero_size(self):
                 with self.assertRaises(ValueError):
                     getattr(mx, op)(a_mx, axis=axis)
 
+    def test_zero_size_sum_prod_all_dtypes(self):
+        # sum and prod have identities, so an empty reduction returns them for
+        # every dtype. The unsigned outputs (uint32, uint64) had no init kernel
+        # on Metal and aborted the process instead.
+        dtypes = [
+            mx.bool_,
+            mx.uint8,
+            mx.uint16,
+            mx.uint32,
+            mx.uint64,
+            mx.int8,
+            mx.int16,
+            mx.int32,
+            mx.int64,
+            mx.float16,
+            mx.float32,
+        ]
+        for dtype in dtypes:
+            a = mx.zeros((0,), dtype=dtype)
+            for op, identity in [("sum", 0), ("prod", 1)]:
+                out = getattr(mx, op)(a)
+                mx.eval(out)
+                self.assertEqual(out.item(), identity, f"{op} of empty {dtype}")
+            # empty because of another axis, reduced over a non-empty one
+            b = mx.zeros((0, 3), dtype=dtype)
+            for op in ["sum", "prod"]:
+                out = getattr(mx, op)(b, axis=-1)
+                mx.eval(out)
+                self.assertEqual(out.shape, (0,))
+
     def test_sum_bool(self):
         x = np.random.uniform(0, 1, size=(10, 10, 10)) > 0.5
         y = mx.array(x)

From c5adcb8de6e3e38ef68c82c5c9102e845cab67bd Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Thu, 1 Oct 2026 11:05:13 +0800
Subject: [PATCH 2/2] nit

---
 python/tests/test_reduce.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py
index 338d7cb7b0..c155ad0a9f 100644
--- a/python/tests/test_reduce.py
+++ b/python/tests/test_reduce.py
@@ -187,7 +187,6 @@ def test_zero_size(self):
                 with self.assertRaises(ValueError):
                     getattr(mx, op)(a_mx, axis=axis)
 
-    def test_zero_size_sum_prod_all_dtypes(self):
         # sum and prod have identities, so an empty reduction returns them for
         # every dtype. The unsigned outputs (uint32, uint64) had no init kernel
         # on Metal and aborted the process instead.