Boolean reduction axis=None logic change

This case now circumvents the call to permute_dims completely
Tests were updated to reflect this change and cover both branches
Also added a test for the axis=() case

Author: ndgrigorian

dpctl/tensor/_utility_functions.py

@@ -11,22 +11,29 @@ def _boolean_reduction(x, axis, keepdims, func):
 
     nd = x.ndim
     if axis is None:
-        axis = tuple(range(nd))
-    if not isinstance(axis, (tuple, list)):
-        axis = (axis,)
-    axis = normalize_axis_tuple(axis, nd, "axis")
+        red_nd = nd
+        # case of a scalar
+        if red_nd == 0:
+            return dpt.astype(x, dpt.bool)
+        x_tmp = x
+        res_shape = tuple()
+        perm = list(range(nd))
+    else:
+        if not isinstance(axis, (tuple, list)):
+            axis = (axis,)
+        axis = normalize_axis_tuple(axis, nd, "axis")
+
+        red_nd = len(axis)
+        # check for axis=()
+        if red_nd == 0:
+            return dpt.astype(x, dpt.bool)
+        perm = [i for i in range(nd) if i not in axis] + list(axis)
+        x_tmp = dpt.permute_dims(x, perm)
+        res_shape = x_tmp.shape[: nd - red_nd]
 
     exec_q = x.sycl_queue
     res_usm_type = x.usm_type
 
-    red_nd = len(axis)
-    if red_nd == 0:
-        return dpt.astype(x, dpt.bool)
-
-    perm = [i for i in range(nd) if i not in axis] + list(axis)
-    x_tmp = dpt.permute_dims(x, perm)
-    res_shape = x_tmp.shape[: nd - red_nd]
-
     wait_list = []
     res_tmp = dpt.empty(
         res_shape,
@@ -59,7 +66,6 @@ def _boolean_reduction(x, axis, keepdims, func):
         inv_perm = sorted(range(nd), key=lambda d: perm[d])
         res = dpt.permute_dims(dpt.reshape(res, res_shape), inv_perm)
     dpctl.SyclEvent.wait_for(wait_list)
-
     return res
 
 

dpctl/tests/test_usm_ndarray_utility_functions.py

@@ -92,6 +92,9 @@ def test_boolean_reduction_keepdims(func):
     assert res.shape == (2, 1, 1, 5, 1)
     assert_array_equal(dpt.asnumpy(res), np.full(res.shape, True))
 
+    res = func(x, axis=None, keepdims=True)
+    assert res.shape == (1,) * x.ndim
+
 
 @pytest.mark.parametrize("func,identity", [(dpt.all, True), (dpt.any, False)])
 def test_boolean_reduction_empty(func, identity):
@@ -119,7 +122,19 @@ def test_boolean_reduction_scalars(func):
     get_queue_or_skip()
 
     x = dpt.ones((), dtype="i4")
-    func(x)
+    assert_equal(dpt.asnumpy(func(x)), True)
+
+    x = dpt.zeros((), dtype="i4")
+    assert_equal(dpt.asnumpy(func(x)), False)
+
+
+@pytest.mark.parametrize("func", [dpt.all, dpt.any])
+def test_boolean_reduction_empty_axis(func):
+    get_queue_or_skip()
+
+    x = dpt.ones((5,), dtype="i4")
+    res = func(x, axis=())
+    assert_array_equal(dpt.asnumpy(res), dpt.asnumpy(x).astype(np.bool_))
 
 
 @pytest.mark.parametrize("func", [dpt.all, dpt.any])