Move shape-related helpers into shape_helpers.py

Author: honno

array_api_tests/array_helpers.py

@@ -1,5 +1,3 @@
-import itertools
-
 from ._array_module import (isnan, all, any, equal, not_equal, logical_and,
                             logical_or, isfinite, greater, less, less_equal,
                             zeros, ones, full, bool, int8, int16, int32,
@@ -23,7 +21,7 @@
            'assert_isinf', 'positive_mathematical_sign',
            'assert_positive_mathematical_sign', 'negative_mathematical_sign',
            'assert_negative_mathematical_sign', 'same_sign',
-           'assert_same_sign', 'ndindex', 'float64',
+           'assert_same_sign', 'float64',
            'asarray', 'full', 'true', 'false', 'isnan']
 
 def zero(shape, dtype):
@@ -319,13 +317,3 @@ def int_to_dtype(x, n, signed):
         if x & highest_bit:
             x = -((~x & mask) + 1)
     return x
-
-def ndindex(shape):
-    """
-    Iterator of n-D indices to an array
-
-    Yields tuples of integers to index every element of an array of shape
-    `shape`. Same as np.ndindex().
-
-    """
-    return itertools.product(*[range(i) for i in shape])

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, Union
+from typing import Any, List, NamedTuple, Optional, Sequence, Tuple, Union
 
 from hypothesis import assume
 from hypothesis.strategies import (SearchStrategy, booleans, composite, floats,
@@ -11,15 +11,15 @@
 
 from . import _array_module as xp
 from . import dtype_helpers as dh
+from . import shape_helpers as sh
 from . import xps
 from ._array_module import _UndefinedStub
 from ._array_module import bool as bool_dtype
 from ._array_module import broadcast_to, eye, float32, float64, full
-from .array_helpers import ndindex
+from .algos import broadcast_shapes
 from .function_stubs import elementwise_functions
 from .pytest_helpers import nargs
 from .typing import Array, DataType, Shape
-from .algos import broadcast_shapes
 
 # Set this to True to not fail tests just because a dtype isn't implemented.
 # If no compatible dtype is implemented for a given test, the test will fail
@@ -208,7 +208,7 @@ def invertible_matrices(draw, dtypes=xps.floating_dtypes(), stack_shapes=shapes(
     assume(xp.all(xp.abs(d) > 0.5))
 
     a = xp.zeros(shape)
-    for j, (idx, i) in enumerate(itertools.product(ndindex(stack_shape), range(n))):
+    for j, (idx, i) in enumerate(itertools.product(sh.ndindex(stack_shape), range(n))):
         a[idx + (i, i)] = d[j]
     return a
 

array_api_tests/meta/test_utils.py

@@ -1,10 +1,8 @@
 import pytest
 
-from .. import array_helpers as ah
+from .. import shape_helpers as sh
 from ..test_creation_functions import frange
-from ..test_manipulation_functions import axis_ndindex
 from ..test_signatures import extension_module
-from ..test_statistical_functions import axes_ndindex
 
 
 def test_extension_module_is_extension():
@@ -34,7 +32,7 @@ def test_frange(r, size, elements):
     [((), [()])],
 )
 def test_ndindex(shape, expected):
-    assert list(ah.ndindex(shape)) == expected
+    assert list(sh.ndindex(shape)) == expected
 
 
 @pytest.mark.parametrize(
@@ -50,7 +48,7 @@ def test_ndindex(shape, expected):
     ],
 )
 def test_axis_ndindex(shape, axis, expected):
-    assert list(axis_ndindex(shape, axis)) == expected
+    assert list(sh.axis_ndindex(shape, axis)) == expected
 
 
 @pytest.mark.parametrize(
@@ -69,4 +67,4 @@ def test_axis_ndindex(shape, axis, expected):
     ],
 )
 def test_axes_ndindex(shape, axes, expected):
-    assert list(axes_ndindex(shape, axes)) == expected
+    assert list(sh.axes_ndindex(shape, axes)) == expected

array_api_tests/shape_helpers.py

@@ -0,0 +1,59 @@
+from itertools import product
+from typing import Iterator, List, Optional, Tuple, Union
+
+from .typing import Shape
+
+__all__ = ["normalise_axis", "ndindex", "axis_ndindex", "axes_ndindex"]
+
+
+def normalise_axis(
+    axis: Optional[Union[int, Tuple[int, ...]]], ndim: int
+) -> Tuple[int, ...]:
+    if axis is None:
+        return tuple(range(ndim))
+    axes = axis if isinstance(axis, tuple) else (axis,)
+    axes = tuple(axis if axis >= 0 else ndim + axis for axis in axes)
+    return axes
+
+
+def ndindex(shape):
+    """Iterator of n-D indices to an array
+
+    Yields tuples of integers to index every element of an array of shape
+    `shape`. Same as np.ndindex().
+    """
+    return product(*[range(i) for i in shape])
+
+
+def axis_ndindex(
+    shape: Shape, axis: int
+) -> Iterator[Tuple[Tuple[Union[int, slice], ...], ...]]:
+    """Generate indices that index all elements in dimensions beyond `axis`"""
+    assert axis >= 0  # sanity check
+    axis_indices = [range(side) for side in shape[:axis]]
+    for _ in range(axis, len(shape)):
+        axis_indices.append([slice(None, None)])
+    yield from product(*axis_indices)
+
+
+def axes_ndindex(shape: Shape, axes: Tuple[int, ...]) -> Iterator[List[Shape]]:
+    """Generate indices that index all elements except in `axes` dimensions"""
+    base_indices = []
+    axes_indices = []
+    for axis, side in enumerate(shape):
+        if axis in axes:
+            base_indices.append([None])
+            axes_indices.append(range(side))
+        else:
+            base_indices.append(range(side))
+            axes_indices.append([None])
+    for base_idx in product(*base_indices):
+        indices = []
+        for idx in product(*axes_indices):
+            idx = list(idx)
+            for axis, side in enumerate(idx):
+                if axis not in axes:
+                    idx[axis] = base_idx[axis]
+            idx = tuple(idx)
+            indices.append(idx)
+        yield list(indices)

array_api_tests/test_elementwise_functions.py

@@ -23,6 +23,7 @@
 from . import dtype_helpers as dh
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
+from . import shape_helpers as sh
 from . import xps
 from .algos import broadcast_shapes
 from .typing import Array, DataType, Param, Scalar
@@ -377,13 +378,13 @@ def test_bitwise_and(
 
         # Compare against the Python & operator.
         if res.dtype == xp.bool:
-            for idx in ah.ndindex(res.shape):
+            for idx in sh.ndindex(res.shape):
                 s_left = bool(_left[idx])
                 s_right = bool(_right[idx])
                 s_res = bool(res[idx])
                 assert (s_left and s_right) == s_res
         else:
-            for idx in ah.ndindex(res.shape):
+            for idx in sh.ndindex(res.shape):
                 s_left = int(_left[idx])
                 s_right = int(_right[idx])
                 s_res = int(res[idx])
@@ -427,7 +428,7 @@ def test_bitwise_left_shift(
         _right = xp.broadcast_to(right, shape)
 
         # Compare against the Python << operator.
-        for idx in ah.ndindex(res.shape):
+        for idx in sh.ndindex(res.shape):
             s_left = int(_left[idx])
             s_right = int(_right[idx])
             s_res = int(res[idx])
@@ -452,12 +453,12 @@ def test_bitwise_invert(func_name, func, strat, data):
     ph.assert_shape(func_name, out.shape, x.shape)
     # Compare against the Python ~ operator.
     if out.dtype == xp.bool:
-        for idx in ah.ndindex(out.shape):
+        for idx in sh.ndindex(out.shape):
             s_x = bool(x[idx])
             s_out = bool(out[idx])
             assert (not s_x) == s_out
     else:
-        for idx in ah.ndindex(out.shape):
+        for idx in sh.ndindex(out.shape):
             s_x = int(x[idx])
             s_out = int(out[idx])
             s_invert = ah.int_to_dtype(
@@ -495,13 +496,13 @@ def test_bitwise_or(
 
         # Compare against the Python | operator.
         if res.dtype == xp.bool:
-            for idx in ah.ndindex(res.shape):
+            for idx in sh.ndindex(res.shape):
                 s_left = bool(_left[idx])
                 s_right = bool(_right[idx])
                 s_res = bool(res[idx])
                 assert (s_left or s_right) == s_res
         else:
-            for idx in ah.ndindex(res.shape):
+            for idx in sh.ndindex(res.shape):
                 s_left = int(_left[idx])
                 s_right = int(_right[idx])
                 s_res = int(res[idx])
@@ -547,7 +548,7 @@ def test_bitwise_right_shift(
         _right = xp.broadcast_to(right, shape)
 
         # Compare against the Python >> operator.
-        for idx in ah.ndindex(res.shape):
+        for idx in sh.ndindex(res.shape):
             s_left = int(_left[idx])
             s_right = int(_right[idx])
             s_res = int(res[idx])
@@ -586,13 +587,13 @@ def test_bitwise_xor(
 
         # Compare against the Python ^ operator.
         if res.dtype == xp.bool:
-            for idx in ah.ndindex(res.shape):
+            for idx in sh.ndindex(res.shape):
                 s_left = bool(_left[idx])
                 s_right = bool(_right[idx])
                 s_res = bool(res[idx])
                 assert (s_left ^ s_right) == s_res
         else:
-            for idx in ah.ndindex(res.shape):
+            for idx in sh.ndindex(res.shape):
                 s_left = int(_left[idx])
                 s_right = int(_right[idx])
                 s_res = int(res[idx])
@@ -721,7 +722,7 @@ def test_equal(
         _right = ah.asarray(_right, dtype=promoted_dtype)
 
         scalar_type = dh.get_scalar_type(promoted_dtype)
-        for idx in ah.ndindex(shape):
+        for idx in sh.ndindex(shape):
             x1_idx = _left[idx]
             x2_idx = _right[idx]
             out_idx = out[idx]
@@ -846,7 +847,7 @@ def test_greater(
         _right = ah.asarray(_right, dtype=promoted_dtype)
 
         scalar_type = dh.get_scalar_type(promoted_dtype)
-        for idx in ah.ndindex(shape):
+        for idx in sh.ndindex(shape):
             out_idx = out[idx]
             x1_idx = _left[idx]
             x2_idx = _right[idx]
@@ -887,7 +888,7 @@ def test_greater_equal(
         _right = ah.asarray(_right, dtype=promoted_dtype)
 
         scalar_type = dh.get_scalar_type(promoted_dtype)
-        for idx in ah.ndindex(shape):
+        for idx in sh.ndindex(shape):
             out_idx = out[idx]
             x1_idx = _left[idx]
             x2_idx = _right[idx]
@@ -907,7 +908,7 @@ def test_isfinite(x):
 
     # Test the exact value by comparing to the math version
     if dh.is_float_dtype(x.dtype):
-        for idx in ah.ndindex(x.shape):
+        for idx in sh.ndindex(x.shape):
             s = float(x[idx])
             assert bool(res[idx]) == math.isfinite(s)
 
@@ -925,7 +926,7 @@ def test_isinf(x):
 
     # Test the exact value by comparing to the math version
     if dh.is_float_dtype(x.dtype):
-        for idx in ah.ndindex(x.shape):
+        for idx in sh.ndindex(x.shape):
             s = float(x[idx])
             assert bool(res[idx]) == math.isinf(s)
 
@@ -943,7 +944,7 @@ def test_isnan(x):
 
     # Test the exact value by comparing to the math version
     if dh.is_float_dtype(x.dtype):
-        for idx in ah.ndindex(x.shape):
+        for idx in sh.ndindex(x.shape):
             s = float(x[idx])
             assert bool(res[idx]) == math.isnan(s)
 
@@ -979,7 +980,7 @@ def test_less(
         _right = ah.asarray(_right, dtype=promoted_dtype)
 
         scalar_type = dh.get_scalar_type(promoted_dtype)
-        for idx in ah.ndindex(shape):
+        for idx in sh.ndindex(shape):
             x1_idx = _left[idx]
             x2_idx = _right[idx]
             out_idx = out[idx]
@@ -1020,7 +1021,7 @@ def test_less_equal(
         _right = ah.asarray(_right, dtype=promoted_dtype)
 
         scalar_type = dh.get_scalar_type(promoted_dtype)
-        for idx in ah.ndindex(shape):
+        for idx in sh.ndindex(shape):
             x1_idx = _left[idx]
             x2_idx = _right[idx]
             out_idx = out[idx]
@@ -1100,15 +1101,15 @@ def test_logical_and(x1, x2):
     _x1 = xp.broadcast_to(x1, shape)
     _x2 = xp.broadcast_to(x2, shape)
 
-    for idx in ah.ndindex(shape):
+    for idx in sh.ndindex(shape):
         assert out[idx] == (bool(_x1[idx]) and bool(_x2[idx]))
 
 
 @given(xps.arrays(dtype=xp.bool, shape=hh.shapes()))
 def test_logical_not(x):
     out = ah.logical_not(x)
     ph.assert_shape("logical_not", out.shape, x.shape)
-    for idx in ah.ndindex(x.shape):
+    for idx in sh.ndindex(x.shape):
         assert out[idx] == (not bool(x[idx]))
 
 
@@ -1122,7 +1123,7 @@ def test_logical_or(x1, x2):
     _x1 = xp.broadcast_to(x1, shape)
     _x2 = xp.broadcast_to(x2, shape)
 
-    for idx in ah.ndindex(shape):
+    for idx in sh.ndindex(shape):
         assert out[idx] == (bool(_x1[idx]) or bool(_x2[idx]))
 
 
@@ -1136,7 +1137,7 @@ def test_logical_xor(x1, x2):
     _x1 = xp.broadcast_to(x1, shape)
     _x2 = xp.broadcast_to(x2, shape)
 
-    for idx in ah.ndindex(shape):
+    for idx in sh.ndindex(shape):
         assert out[idx] == (bool(_x1[idx]) ^ bool(_x2[idx]))
 
 
@@ -1225,7 +1226,7 @@ def test_not_equal(
         _right = ah.asarray(_right, dtype=promoted_dtype)
 
         scalar_type = dh.get_scalar_type(promoted_dtype)
-        for idx in ah.ndindex(shape):
+        for idx in sh.ndindex(shape):
             out_idx = out[idx]
             x1_idx = _left[idx]
             x2_idx = _right[idx]