Docstrings for axes helpers

honno · honno · commit c1ca9fb6e4ea · 2021-12-08T11:15:54.000Z
diff --git a/array_api_tests/test_manipulation_functions.py b/array_api_tests/test_manipulation_functions.py
@@ -13,7 +13,7 @@
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
 from . import xps
-from .test_statistical_functions import axes_ndindex, normalise_axis  # TODO: Move
+from .test_statistical_functions import axes_ndindex, normalise_axis
 from .typing import Array, Shape
 
 MAX_SIDE = hh.MAX_ARRAY_SIZE // 64
@@ -32,6 +32,7 @@ def shared_shapes(*args, **kwargs) -> st.SearchStrategy[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)):
diff --git a/array_api_tests/test_statistical_functions.py b/array_api_tests/test_statistical_functions.py
@@ -39,18 +39,19 @@ def normalise_axis(
 
 
 def axes_ndindex(shape: Shape, axes: Tuple[int, ...]) -> Iterator[Tuple[Shape, ...]]:
-    base_iterables = []
-    axes_iterables = []
+    """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_iterables.append([None])
-            axes_iterables.append(range(side))
+            base_indices.append([None])
+            axes_indices.append(range(side))
         else:
-            base_iterables.append(range(side))
-            axes_iterables.append([None])
-    for base_idx in product(*base_iterables):
+            base_indices.append(range(side))
+            axes_indices.append([None])
+    for base_idx in product(*base_indices):
         indices = []
-        for idx in product(*axes_iterables):
+        for idx in product(*axes_indices):
             idx = list(idx)
             for axis, side in enumerate(idx):
                 if axis not in axes:
