Swap {in/out}_shape params in ph.assert_keepdimable_shape()

Also documents it

Author: honno

array_api_tests/pytest_helpers.py

@@ -264,13 +264,28 @@ def assert_result_shape(
 
 def assert_keepdimable_shape(
     func_name: str,
-    out_shape: Shape,
     in_shape: Shape,
+    out_shape: Shape,
     axes: Tuple[int, ...],
     keepdims: bool,
     /,
     **kw,
 ):
+    """
+    Assert the output shape from a keepdimable function is as expected, e.g.
+
+        >>> x = xp.asarray([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
+        >>> out1 = xp.max(x, keepdims=False)
+        >>> out2 = xp.max(x, keepdims=True)
+        >>> assert_keepdimable_shape('max', x.shape, out1.shape, (0, 1), False)
+        >>> assert_keepdimable_shape('max', x.shape, out2.shape, (0, 1), True)
+
+        is equivalent to
+
+        >>> assert out1.shape == ()
+        >>> assert out2.shape == (1, 1)
+
+    """
     if keepdims:
         shape = tuple(1 if axis in axes else side for axis, side in enumerate(in_shape))
     else:

array_api_tests/test_searching_functions.py

@@ -34,7 +34,7 @@ def test_argmax(x, data):
     ph.assert_default_index("argmax", out.dtype)
     axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
-        "argmax", out.shape, x.shape, axes, kw.get("keepdims", False), **kw
+        "argmax", x.shape, out.shape, axes, kw.get("keepdims", False), **kw
     )
     scalar_type = dh.get_scalar_type(x.dtype)
     for indices, out_idx in zip(sh.axes_ndindex(x.shape, axes), sh.ndindex(out.shape)):
@@ -69,7 +69,7 @@ def test_argmin(x, data):
     ph.assert_default_index("argmin", out.dtype)
     axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
-        "argmin", out.shape, x.shape, axes, kw.get("keepdims", False), **kw
+        "argmin", x.shape, out.shape, axes, kw.get("keepdims", False), **kw
     )
     scalar_type = dh.get_scalar_type(x.dtype)
     for indices, out_idx in zip(sh.axes_ndindex(x.shape, axes), sh.ndindex(out.shape)):

array_api_tests/test_statistical_functions.py

@@ -38,7 +38,7 @@ def test_max(x, data):
     ph.assert_dtype("max", x.dtype, out.dtype)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
-        "max", out.shape, x.shape, _axes, kw.get("keepdims", False), **kw
+        "max", x.shape, out.shape, _axes, kw.get("keepdims", False), **kw
     )
     scalar_type = dh.get_scalar_type(out.dtype)
     for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
@@ -67,7 +67,7 @@ def test_mean(x, data):
     ph.assert_dtype("mean", x.dtype, out.dtype)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
-        "mean", out.shape, x.shape, _axes, kw.get("keepdims", False), **kw
+        "mean", x.shape, out.shape, _axes, kw.get("keepdims", False), **kw
     )
     # Values testing mean is too finicky
 
@@ -88,7 +88,7 @@ def test_min(x, data):
     ph.assert_dtype("min", x.dtype, out.dtype)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
-        "min", out.shape, x.shape, _axes, kw.get("keepdims", False), **kw
+        "min", x.shape, out.shape, _axes, kw.get("keepdims", False), **kw
     )
     scalar_type = dh.get_scalar_type(out.dtype)
     for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
@@ -147,7 +147,7 @@ def test_prod(x, data):
     ph.assert_dtype("prod", x.dtype, out.dtype, _dtype)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
-        "prod", out.shape, x.shape, _axes, kw.get("keepdims", False), **kw
+        "prod", x.shape, out.shape, _axes, kw.get("keepdims", False), **kw
     )
     scalar_type = dh.get_scalar_type(out.dtype)
     for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
@@ -194,7 +194,7 @@ def test_std(x, data):
 
     ph.assert_dtype("std", x.dtype, out.dtype)
     ph.assert_keepdimable_shape(
-        "std", out.shape, x.shape, _axes, kw.get("keepdims", False), **kw
+        "std", x.shape, out.shape, _axes, kw.get("keepdims", False), **kw
     )
     # We can't easily test the result(s) as standard deviation methods vary a lot
 
@@ -245,7 +245,7 @@ def test_sum(x, data):
     ph.assert_dtype("sum", x.dtype, out.dtype, _dtype)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
-        "sum", out.shape, x.shape, _axes, kw.get("keepdims", False), **kw
+        "sum", x.shape, out.shape, _axes, kw.get("keepdims", False), **kw
     )
     scalar_type = dh.get_scalar_type(out.dtype)
     for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
@@ -292,6 +292,6 @@ def test_var(x, data):
 
     ph.assert_dtype("var", x.dtype, out.dtype)
     ph.assert_keepdimable_shape(
-        "var", out.shape, x.shape, _axes, kw.get("keepdims", False), **kw
+        "var", x.shape, out.shape, _axes, kw.get("keepdims", False), **kw
     )
     # We can't easily test the result(s) as variance methods vary a lot

array_api_tests/test_utility_functions.py

@@ -24,7 +24,7 @@ def test_all(x, data):
     ph.assert_dtype("all", x.dtype, out.dtype, xp.bool)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
-        "all", out.shape, x.shape, _axes, kw.get("keepdims", False), **kw
+        "all", x.shape, out.shape, _axes, kw.get("keepdims", False), **kw
     )
     scalar_type = dh.get_scalar_type(x.dtype)
     for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
@@ -49,7 +49,7 @@ def test_any(x, data):
     ph.assert_dtype("any", x.dtype, out.dtype, xp.bool)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
-        "any", out.shape, x.shape, _axes, kw.get("keepdims", False), **kw
+        "any", x.shape, out.shape, _axes, kw.get("keepdims", False), **kw
     )
     scalar_type = dh.get_scalar_type(x.dtype)
     for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):