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..c155ad0a9f 100644 --- a/python/tests/test_reduce.py +++ b/python/tests/test_reduce.py @@ -187,6 +187,35 @@ def test_zero_size(self): with self.assertRaises(ValueError): getattr(mx, op)(a_mx, axis=axis) + # 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)