Move axes() strategy to hypothesis_helpers.py

Author: honno

array_api_tests/hypothesis_helpers.py

@@ -2,7 +2,7 @@
 from functools import reduce
 from math import sqrt
 from operator import mul
-from typing import Any, List, NamedTuple, Optional, Tuple, Sequence
+from typing import Any, List, NamedTuple, Optional, Tuple, Sequence, Union
 
 from hypothesis import assume
 from hypothesis.strategies import (SearchStrategy, booleans, composite, floats,
@@ -399,3 +399,12 @@ def specified_kwargs(draw, *keys_values_defaults: KVD):
         if value is not default or draw(booleans()):
             kw[keyword] = value
     return kw
+
+
+def axes(ndim: int) -> SearchStrategy[Optional[Union[int, Shape]]]:
+    """Generate valid arguments for some axis keywords"""
+    axes_strats = [none()]
+    if ndim != 0:
+        axes_strats.append(integers(-ndim, ndim - 1))
+        axes_strats.append(xps.valid_tuple_axes(ndim))
+    return one_of(axes_strats)

array_api_tests/pytest_helpers.py

@@ -1,4 +1,4 @@
-from array_api_tests.algos import broadcast_shapes
+from .algos import broadcast_shapes
 import math
 from inspect import getfullargspec
 from typing import Any, Dict, Optional, Tuple, Union

array_api_tests/test_searching_functions.py

@@ -1,22 +1,21 @@
 from hypothesis import given
 from hypothesis import strategies as st
 
-from array_api_tests.algos import broadcast_shapes
-from array_api_tests.test_manipulation_functions import assert_equals as assert_equals_
-from array_api_tests.test_statistical_functions import (
-    assert_equals,
-    assert_keepdimable_shape,
-    axes_ndindex,
-    normalise_axis,
-)
-from array_api_tests.typing import DataType
-
 from . import _array_module as xp
 from . import array_helpers as ah
 from . import dtype_helpers as dh
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
 from . import xps
+from .algos import broadcast_shapes
+from .test_manipulation_functions import assert_equals as assert_equals_
+from .test_statistical_functions import (
+    assert_equals,
+    assert_keepdimable_shape,
+    axes_ndindex,
+    normalise_axis,
+)
+from .typing import DataType
 
 
 def assert_default_index(func_name: str, dtype: DataType, repr_name="out.dtype"):

array_api_tests/test_statistical_functions.py

@@ -15,14 +15,6 @@
 from .typing import DataType, Scalar, ScalarType, Shape
 
 
-def axes(ndim: int) -> st.SearchStrategy[Optional[Union[int, Shape]]]:
-    axes_strats = [st.none()]
-    if ndim != 0:
-        axes_strats.append(st.integers(-ndim, ndim - 1))
-        axes_strats.append(xps.valid_tuple_axes(ndim))
-    return st.one_of(axes_strats)
-
-
 def kwarg_dtypes(dtype: DataType) -> st.SearchStrategy[Optional[DataType]]:
     dtypes = [d2 for d1, d2 in dh.promotion_table if d1 == dtype]
     return st.none() | st.sampled_from(dtypes)
@@ -108,7 +100,7 @@ def assert_equals(
     data=st.data(),
 )
 def test_max(x, data):
-    kw = data.draw(hh.kwargs(axis=axes(x.ndim), keepdims=st.booleans()), label="kw")
+    kw = data.draw(hh.kwargs(axis=hh.axes(x.ndim), keepdims=st.booleans()), label="kw")
 
     out = xp.max(x, **kw)
 
@@ -137,7 +129,7 @@ def test_max(x, data):
     data=st.data(),
 )
 def test_mean(x, data):
-    kw = data.draw(hh.kwargs(axis=axes(x.ndim), keepdims=st.booleans()), label="kw")
+    kw = data.draw(hh.kwargs(axis=hh.axes(x.ndim), keepdims=st.booleans()), label="kw")
 
     out = xp.mean(x, **kw)
 
@@ -166,7 +158,7 @@ def test_mean(x, data):
     data=st.data(),
 )
 def test_min(x, data):
-    kw = data.draw(hh.kwargs(axis=axes(x.ndim), keepdims=st.booleans()), label="kw")
+    kw = data.draw(hh.kwargs(axis=hh.axes(x.ndim), keepdims=st.booleans()), label="kw")
 
     out = xp.min(x, **kw)
 
@@ -197,7 +189,7 @@ def test_min(x, data):
 def test_prod(x, data):
     kw = data.draw(
         hh.kwargs(
-            axis=axes(x.ndim),
+            axis=hh.axes(x.ndim),
             dtype=kwarg_dtypes(x.dtype),
             keepdims=st.booleans(),
         ),
@@ -258,7 +250,7 @@ def test_prod(x, data):
     data=st.data(),
 )
 def test_std(x, data):
-    axis = data.draw(axes(x.ndim), label="axis")
+    axis = data.draw(hh.axes(x.ndim), label="axis")
     _axes = normalise_axis(axis, x.ndim)
     N = sum(side for axis, side in enumerate(x.shape) if axis not in _axes)
     correction = data.draw(
@@ -295,7 +287,7 @@ def test_std(x, data):
 def test_sum(x, data):
     kw = data.draw(
         hh.kwargs(
-            axis=axes(x.ndim),
+            axis=hh.axes(x.ndim),
             dtype=kwarg_dtypes(x.dtype),
             keepdims=st.booleans(),
         ),
@@ -356,7 +348,7 @@ def test_sum(x, data):
     data=st.data(),
 )
 def test_var(x, data):
-    axis = data.draw(axes(x.ndim), label="axis")
+    axis = data.draw(hh.axes(x.ndim), label="axis")
     _axes = normalise_axis(axis, x.ndim)
     N = sum(side for axis, side in enumerate(x.shape) if axis not in _axes)
     correction = data.draw(

array_api_tests/test_utility_functions.py

@@ -10,7 +10,6 @@
 from .test_statistical_functions import (
     assert_equals,
     assert_keepdimable_shape,
-    axes,
     axes_ndindex,
     normalise_axis,
 )
@@ -21,7 +20,7 @@
     data=st.data(),
 )
 def test_all(x, data):
-    kw = data.draw(hh.kwargs(axis=axes(x.ndim), keepdims=st.booleans()), label="kw")
+    kw = data.draw(hh.kwargs(axis=hh.axes(x.ndim), keepdims=st.booleans()), label="kw")
 
     out = xp.all(x, **kw)
 
@@ -46,7 +45,7 @@ def test_all(x, data):
     data=st.data(),
 )
 def test_any(x, data):
-    kw = data.draw(hh.kwargs(axis=axes(x.ndim), keepdims=st.booleans()), label="kw")
+    kw = data.draw(hh.kwargs(axis=hh.axes(x.ndim), keepdims=st.booleans()), label="kw")
 
     out = xp.any(x, **kw)