Remove unused conditional axis_strat definition in test_arg*

data-apis · Feb 11, 2022 · 511929c · 511929c
1 parent aae2165
commit 511929c
Showing 1 changed file with 14 additions and 8 deletions.
diff --git a/array_api_tests/test_searching_functions.py b/array_api_tests/test_searching_functions.py
@@ -21,10 +21,13 @@
     data=st.data(),
 )
 def test_argmax(x, data):
-    axis_strat = st.none()
-    if x.ndim > 0:
-        axis_strat |= st.integers(-x.ndim, max(x.ndim - 1, 0))
-    kw = data.draw(hh.kwargs(axis=axis_strat, keepdims=st.booleans()), label="kw")
+    kw = data.draw(
+        hh.kwargs(
+            axis=st.none() | st.integers(-x.ndim, max(x.ndim - 1, 0)),
+            keepdims=st.booleans(),
+        ),
+        label="kw",
+    )
 
     out = xp.argmax(x, **kw)
 
@@ -53,10 +56,13 @@ def test_argmax(x, data):
     data=st.data(),
 )
 def test_argmin(x, data):
-    axis_strat = st.none()
-    if x.ndim > 0:
-        axis_strat |= st.integers(-x.ndim, max(x.ndim - 1, 0))
-    kw = data.draw(hh.kwargs(axis=axis_strat, keepdims=st.booleans()), label="kw")
+    kw = data.draw(
+        hh.kwargs(
+            axis=st.none() | st.integers(-x.ndim, max(x.ndim - 1, 0)),
+            keepdims=st.booleans(),
+        ),
+        label="kw",
+    )
 
     out = xp.argmin(x, **kw)