Fix matmul memory test by passing dtype properly (#358)

cubed-dev · Jan 25, 2024 · 2e28072 · 2e28072
1 parent 929f52c
commit 2e28072
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/cubed/array_api/linear_algebra_functions.py b/cubed/array_api/linear_algebra_functions.py
@@ -68,7 +68,10 @@ def _sum_wo_cat(a, axis=None, dtype=None):
     if a.shape[axis] == 1:
         return squeeze(a, axis)
 
-    return reduction(a, _chunk_sum, axis=axis, dtype=dtype)
+    extra_func_kwargs = dict(dtype=dtype)
+    return reduction(
+        a, _chunk_sum, axis=axis, dtype=dtype, extra_func_kwargs=extra_func_kwargs
+    )
 
 
 def _chunk_sum(a, axis=None, dtype=None, keepdims=None):