Bug/999 Fix keepdim in any/all

heat/core/logical.py

@@ -91,6 +91,9 @@ def all(
     def local_all(t, *args, **kwargs):
         return torch.all(t != 0, *args, **kwargs)
 
+    if keepdim and axis is None:
+        axis = tuple(range(x.ndim))
+
     return _operations.__reduce_op(
         x, local_all, MPI.LAND, axis=axis, out=out, neutral=1, keepdim=keepdim
     )
@@ -196,6 +199,9 @@ def any(
     def local_any(t, *args, **kwargs):
         return torch.any(t != 0, *args, **kwargs)
 
+    if keepdim and axis is None:
+        axis = tuple(range(x.ndim))
+
     return _operations.__reduce_op(
         x, local_any, MPI.LOR, axis=axis, out=out, neutral=0, keepdim=keepdim
     )

heat/core/tests/test_logical.py

@@ -140,6 +140,32 @@ def test_all(self):
         out_noaxis = ht.zeros((1, 2, 3, 5), split=1)
         ht.all(ones_noaxis_split_axis_neg, axis=-2, out=out_noaxis)
 
+        # test keepdim
+        ones_2d = ht.ones((1, 1))
+        self.assertEqual(ones_2d.all(keepdim=True).shape, ones_2d.shape)
+
+        ones_2d_split = ht.ones((2, 2), split=0)
+        keepdim_is_one = ones_2d_split.all(keepdim=True)
+        self.assertEqual(keepdim_is_one.shape, (1, 1))
+        self.assertEqual(keepdim_is_one.split, None)
+        keepdim_is_one = ones_2d_split.all(axis=0, keepdim=True)
+        self.assertEqual(keepdim_is_one.shape, (1, 2))
+        self.assertEqual(keepdim_is_one.split, None)
+        keepdim_is_one = ones_2d_split.all(axis=1, keepdim=True)
+        self.assertEqual(keepdim_is_one.shape, (2, 1))
+        self.assertEqual(keepdim_is_one.split, 0)
+
+        ones_2d_split = ht.ones((2, 2), split=1)
+        keepdim_is_one = ones_2d_split.all(keepdim=True)
+        self.assertEqual(keepdim_is_one.shape, (1, 1))
+        self.assertEqual(keepdim_is_one.split, None)
+        keepdim_is_one = ones_2d_split.all(axis=0, keepdim=True)
+        self.assertEqual(keepdim_is_one.shape, (1, 2))
+        self.assertEqual(keepdim_is_one.split, 1)
+        keepdim_is_one = ones_2d_split.all(axis=1, keepdim=True)
+        self.assertEqual(keepdim_is_one.shape, (2, 1))
+        self.assertEqual(keepdim_is_one.split, None)
+
         # exceptions
         with self.assertRaises(ValueError):
             ht.ones(array_len).all(axis=1)
@@ -212,6 +238,32 @@ def test_any(self):
         self.assertEqual(any_tensor.dtype, ht.bool)
         self.assertTrue(ht.equal(any_tensor, res))
 
+        # test keepdim
+        ones_2d = ht.ones((1, 1))
+        self.assertEqual(ones_2d.any(keepdim=True).shape, ones_2d.shape)
+
+        ones_2d_split = ht.ones((2, 2), split=0)
+        keepdim_any = ones_2d_split.any(keepdim=True)
+        self.assertEqual(keepdim_any.shape, (1, 1))
+        self.assertEqual(keepdim_any.split, None)
+        keepdim_any = ones_2d_split.any(axis=0, keepdim=True)
+        self.assertEqual(keepdim_any.shape, (1, 2))
+        self.assertEqual(keepdim_any.split, None)
+        keepdim_any = ones_2d_split.any(axis=1, keepdim=True)
+        self.assertEqual(keepdim_any.shape, (2, 1))
+        self.assertEqual(keepdim_any.split, 0)
+
+        ones_2d_split = ht.ones((2, 2), split=1)
+        keepdim_any = ones_2d_split.any(keepdim=True)
+        self.assertEqual(keepdim_any.shape, (1, 1))
+        self.assertEqual(keepdim_any.split, None)
+        keepdim_any = ones_2d_split.any(axis=0, keepdim=True)
+        self.assertEqual(keepdim_any.shape, (1, 2))
+        self.assertEqual(keepdim_any.split, 1)
+        keepdim_any = ones_2d_split.any(axis=1, keepdim=True)
+        self.assertEqual(keepdim_any.shape, (2, 1))
+        self.assertEqual(keepdim_any.split, None)
+
     def test_isclose(self):
         size = ht.communication.MPI_WORLD.size
         a = ht.float32([[2, 2], [2, 2]])