Cover axis in test_concat

Author: honno

array_api_tests/meta/test_utils.py

@@ -1,6 +1,8 @@
 import pytest
 
+from .. import array_helpers as ah
 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
 
@@ -27,6 +29,30 @@ def test_frange(r, size, elements):
     assert list(r) == elements
 
 
+@pytest.mark.parametrize(
+    "shape, expected",
+    [((), [()])],
+)
+def test_ndindex(shape, expected):
+    assert list(ah.ndindex(shape)) == expected
+
+
+@pytest.mark.parametrize(
+    "shape, axis, expected",
+    [
+        ((1,), 0, [(slice(None, None),)]),
+        ((1, 2), 0, [(slice(None, None), slice(None, None))]),
+        (
+            (2, 4),
+            1,
+            [(0, slice(None, None)), (1, slice(None, None))],
+        ),
+    ],
+)
+def test_axis_ndindex(shape, axis, expected):
+    assert list(axis_ndindex(shape, axis)) == expected
+
+
 @pytest.mark.parametrize(
     "shape, axes, expected",
     [

array_api_tests/test_manipulation_functions.py

@@ -1,6 +1,7 @@
 import math
 from collections import deque
-from typing import Iterable, Union
+from itertools import product
+from typing import Iterable, Iterator, Tuple, Union
 
 from hypothesis import assume, given
 from hypothesis import strategies as st
@@ -43,6 +44,28 @@ def assert_array_ndindex(
             assert out[out_idx] == x[x_idx], msg
 
 
+def axis_ndindex(
+    shape: Shape, axis: int
+) -> Iterator[Tuple[Tuple[Union[int, slice], ...], ...]]:
+    iterables = [range(side) for side in shape[:axis]]
+    for _ in range(len(shape[axis:])):
+        iterables.append([slice(None, None)])
+    yield from product(*iterables)
+
+
+def assert_equals(
+    func_name: str, x_repr: str, x_val: Array, out_repr: str, out_val: Array, **kw
+):
+    msg = (
+        f"{out_repr}={out_val}, should be {x_repr}={x_val} "
+        f"[{func_name}({ph.fmt_kw(kw)})]"
+    )
+    if dh.is_float_dtype(out_val.dtype) and xp.isnan(out_val):
+        assert xp.isnan(x_val), msg
+    else:
+        assert out_val == out_val, msg
+
+
 @st.composite
 def concat_shapes(draw, shape, axis):
     shape = list(shape)
@@ -85,21 +108,35 @@ def test_concat(dtypes, kw, data):
         shape = tuple(shape)
     ph.assert_result_shape("concat", shapes, out.shape, shape, **kw)
 
-    # TODO: adjust indices with nonzero axis
-    if axis is None or axis == 0:
-        out_indices = ah.ndindex(out.shape)
-        for i, x in enumerate(arrays, 1):
-            msg_suffix = f" [concat({ph.fmt_kw(kw)})]\nx{i}={x!r}\n{out=}"
+    if axis is None:
+        out_indices = (i for i in range(out.size))
+        for x_num, x in enumerate(arrays, 1):
             for x_idx in ah.ndindex(x.shape):
-                out_idx = next(out_indices)
-                msg = (
-                    f"out[{out_idx}]={out[out_idx]}, should be x{i}[{x_idx}]={x[x_idx]}"
+                out_i = next(out_indices)
+                assert_equals(
+                    "concat",
+                    f"x{x_num}[{x_idx}]",
+                    x[x_idx],
+                    f"out[{out_i}]",
+                    out[out_i],
+                    **kw,
                 )
-                msg += msg_suffix
-                if dh.is_float_dtype(x.dtype) and xp.isnan(x[x_idx]):
-                    assert xp.isnan(out[out_idx]), msg
-                else:
-                    assert out[out_idx] == x[x_idx], msg
+    else:
+        out_indices = ah.ndindex(out.shape)
+        for idx in axis_ndindex(shape, axis):
+            f_idx = ", ".join(str(i) if isinstance(i, int) else ":" for i in idx)
+            for x_num, x in enumerate(arrays, 1):
+                indexed_x = x[idx]
+                for x_idx in ah.ndindex(indexed_x.shape):
+                    out_idx = next(out_indices)
+                    assert_equals(
+                        "concat",
+                        f"x{x_num}[{f_idx}][{x_idx}]",
+                        indexed_x[x_idx],
+                        f"out[{out_idx}]",
+                        out[out_idx],
+                        **kw,
+                    )
 
 
 @given(

array_api_tests/test_statistical_functions.py

@@ -42,11 +42,11 @@ def axes_ndindex(shape: Shape, axes: Tuple[int, ...]) -> Iterator[Tuple[Shape, .
     axes_iterables = []
     for axis, side in enumerate(shape):
         if axis in axes:
-            base_iterables.append((None,))
+            base_iterables.append([None])
             axes_iterables.append(range(side))
         else:
             base_iterables.append(range(side))
-            axes_iterables.append((None,))
+            axes_iterables.append([None])
     for base_idx in product(*base_iterables):
         indices = []
         for idx in product(*axes_iterables):