Merge pull request #63 from Quansight-Labs/argminmax_bool

MAINT: work around arg{min,max} not implemented for booleans

Author: ev-br

torch_np/_detail/_reductions.py

@@ -36,12 +36,22 @@ def count_nonzero(a, axis=None):
 
 def argmax(tensor, axis=None):
     axis = _util.allow_only_single_axis(axis)
+
+    if tensor.dtype == torch.bool:
+        # RuntimeError: "argmax_cpu" not implemented for 'Bool'
+        tensor = tensor.to(torch.uint8)
+
     tensor = torch.argmax(tensor, axis)
     return tensor
 
 
 def argmin(tensor, axis=None):
     axis = _util.allow_only_single_axis(axis)
+
+    if tensor.dtype == torch.bool:
+        # RuntimeError: "argmin_cpu" not implemented for 'Bool'
+        tensor = tensor.to(torch.uint8)
+
     tensor = torch.argmin(tensor, axis)
     return tensor
 

torch_np/tests/test_ndarray_methods.py

@@ -407,10 +407,8 @@ def test_combinations(self, data):
         #      with suppress_warnings() as sup:
         #          sup.filter(RuntimeWarning,
         #                      "invalid value encountered in reduce")
-        if np.asarray(arr).dtype.kind in "cb":
-            pytest.xfail(
-                reason="'max_values_cpu' not implemented for 'ComplexDouble', 'Bool'"
-            )
+        if np.asarray(arr).dtype.kind in "c":
+            pytest.xfail(reason="'max_values_cpu' not implemented for 'ComplexDouble'")
 
         val = np.max(arr)
 
@@ -511,10 +509,8 @@ class TestArgmin:
     def test_combinations(self, data):
         arr, pos = data
 
-        if np.asarray(arr).dtype.kind in "cb":
-            pytest.xfail(
-                reason="'min_values_cpu' not implemented for 'ComplexDouble', 'Bool'"
-            )
+        if np.asarray(arr).dtype.kind in "c":
+            pytest.xfail(reason="'min_values_cpu' not implemented for 'ComplexDouble'")
 
         #        with suppress_warnings() as sup:
         #            sup.filter(RuntimeWarning, "invalid value encountered in reduce")