ENH: add "repro snippets" to test_searching_functions.py

Author: ev-br

array_api_tests/test_searching_functions.py

@@ -33,23 +33,28 @@ def test_argmax(x, data):
     )
     keepdims = kw.get("keepdims", False)
 
-    out = xp.argmax(x, **kw)
+    repro_snippet = ph.format_snippet(f"xp.argmax({x!r}, **kw) with {kw = }")
+    try:
+        out = xp.argmax(x, **kw)
 
-    ph.assert_default_index("argmax", out.dtype)
-    axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
-    ph.assert_keepdimable_shape(
-        "argmax", in_shape=x.shape, out_shape=out.shape, axes=axes, keepdims=keepdims, kw=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)):
-        max_i = int(out[out_idx])
-        elements = []
-        for idx in indices:
-            s = scalar_type(x[idx])
-            elements.append(s)
-        expected = max(range(len(elements)), key=elements.__getitem__)
-        ph.assert_scalar_equals("argmax", type_=int, idx=out_idx, out=max_i,
-                                expected=expected, kw=kw)
+        ph.assert_default_index("argmax", out.dtype)
+        axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
+        ph.assert_keepdimable_shape(
+            "argmax", in_shape=x.shape, out_shape=out.shape, axes=axes, keepdims=keepdims, kw=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)):
+            max_i = int(out[out_idx])
+            elements = []
+            for idx in indices:
+                s = scalar_type(x[idx])
+                elements.append(s)
+            expected = max(range(len(elements)), key=elements.__getitem__)
+            ph.assert_scalar_equals("argmax", type_=int, idx=out_idx, out=max_i,
+                                    expected=expected, kw=kw)
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise
 
 
 @given(
@@ -70,22 +75,27 @@ def test_argmin(x, data):
     )
     keepdims = kw.get("keepdims", False)
 
-    out = xp.argmin(x, **kw)
+    repro_snippet = ph.format_snippet(f"xp.argmin({x!r}, **kw) with {kw = }")
+    try:
+        out = xp.argmin(x, **kw)
 
-    ph.assert_default_index("argmin", out.dtype)
-    axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
-    ph.assert_keepdimable_shape(
-        "argmin", in_shape=x.shape, out_shape=out.shape, axes=axes, keepdims=keepdims, kw=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)):
-        min_i = int(out[out_idx])
-        elements = []
-        for idx in indices:
-            s = scalar_type(x[idx])
-            elements.append(s)
-        expected = min(range(len(elements)), key=elements.__getitem__)
-        ph.assert_scalar_equals("argmin", type_=int, idx=out_idx, out=min_i, expected=expected)
+        ph.assert_default_index("argmin", out.dtype)
+        axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
+        ph.assert_keepdimable_shape(
+            "argmin", in_shape=x.shape, out_shape=out.shape, axes=axes, keepdims=keepdims, kw=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)):
+            min_i = int(out[out_idx])
+            elements = []
+            for idx in indices:
+                s = scalar_type(x[idx])
+                elements.append(s)
+            expected = min(range(len(elements)), key=elements.__getitem__)
+            ph.assert_scalar_equals("argmin", type_=int, idx=out_idx, out=min_i, expected=expected)
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise
 
 
 # XXX: the strategy for x is problematic on JAX unless JAX_ENABLE_X64 is on
@@ -115,23 +125,28 @@ def test_count_nonzero(x, data):
 
     assume(kw.get("axis", None) != ())  # TODO clarify in the spec
 
-    out = xp.count_nonzero(x, **kw)
+    repro_snippet = ph.format_snippet(f"xp.count_nonzero({x!r}, **kw) with {kw = }")
+    try:
+        out = xp.count_nonzero(x, **kw)
 
-    ph.assert_default_index("count_nonzero", out.dtype)
-    axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
-    ph.assert_keepdimable_shape(
-        "count_nonzero", in_shape=x.shape, out_shape=out.shape, axes=axes, keepdims=keepdims, kw=kw
-    )
-    scalar_type = dh.get_scalar_type(x.dtype)
+        ph.assert_default_index("count_nonzero", out.dtype)
+        axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
+        ph.assert_keepdimable_shape(
+            "count_nonzero", in_shape=x.shape, out_shape=out.shape, axes=axes, keepdims=keepdims, kw=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)):
-        count = int(out[out_idx])
-        elements = []
-        for idx in indices:
-            s = scalar_type(x[idx])
-            elements.append(s)
-        expected = sum(el != 0 for el in elements)
-        ph.assert_scalar_equals("count_nonzero", type_=int, idx=out_idx, out=count, expected=expected)
+        for indices, out_idx in zip(sh.axes_ndindex(x.shape, axes), sh.ndindex(out.shape)):
+            count = int(out[out_idx])
+            elements = []
+            for idx in indices:
+                s = scalar_type(x[idx])
+                elements.append(s)
+            expected = sum(el != 0 for el in elements)
+            ph.assert_scalar_equals("count_nonzero", type_=int, idx=out_idx, out=count, expected=expected)
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise
 
 
 @given(hh.arrays(dtype=hh.all_dtypes, shape=()))
@@ -143,39 +158,44 @@ def test_nonzero_zerodim_error(x):
 @pytest.mark.data_dependent_shapes
 @given(hh.arrays(dtype=hh.all_dtypes, shape=hh.shapes(min_dims=1, min_side=1)))
 def test_nonzero(x):
-    out = xp.nonzero(x)
-    assert len(out) == x.ndim, f"{len(out)=}, but should be {x.ndim=}"
-    out_size = math.prod(out[0].shape)
-    for i in range(len(out)):
-        assert out[i].ndim == 1, f"out[{i}].ndim={x.ndim}, but should be 1"
-        size_at = math.prod(out[i].shape)
-        assert size_at == out_size, (
-            f"prod(out[{i}].shape)={size_at}, "
-            f"but should be prod(out[0].shape)={out_size}"
-        )
-        ph.assert_default_index("nonzero", out[i].dtype, repr_name=f"out[{i}].dtype")
-    indices = []
-    if x.dtype == xp.bool:
-        for idx in sh.ndindex(x.shape):
-            if x[idx]:
-                indices.append(idx)
-    else:
-        for idx in sh.ndindex(x.shape):
-            if x[idx] != 0:
-                indices.append(idx)
-    if x.ndim == 0:
-        assert out_size == len(
-            indices
-        ), f"prod(out[0].shape)={out_size}, but should be {len(indices)}"
-    else:
-        for i in range(out_size):
-            idx = tuple(int(x[i]) for x in out)
-            f_idx = f"Extrapolated index (x[{i}] for x in out)={idx}"
-            f_element = f"x[{idx}]={x[idx]}"
-            assert idx in indices, f"{f_idx} results in {f_element}, a zero element"
-            assert (
-                idx == indices[i]
-            ), f"{f_idx} is in the wrong position, should be {indices.index(idx)}"
+    repro_snippet = ph.format_snippet(f"xp.nonzero({x!r})")
+    try:
+        out = xp.nonzero(x)
+        assert len(out) == x.ndim, f"{len(out)=}, but should be {x.ndim=}"
+        out_size = math.prod(out[0].shape)
+        for i in range(len(out)):
+            assert out[i].ndim == 1, f"out[{i}].ndim={x.ndim}, but should be 1"
+            size_at = math.prod(out[i].shape)
+            assert size_at == out_size, (
+                f"prod(out[{i}].shape)={size_at}, "
+                f"but should be prod(out[0].shape)={out_size}"
+            )
+            ph.assert_default_index("nonzero", out[i].dtype, repr_name=f"out[{i}].dtype")
+        indices = []
+        if x.dtype == xp.bool:
+            for idx in sh.ndindex(x.shape):
+                if x[idx]:
+                    indices.append(idx)
+        else:
+            for idx in sh.ndindex(x.shape):
+                if x[idx] != 0:
+                    indices.append(idx)
+        if x.ndim == 0:
+            assert out_size == len(
+                indices
+            ), f"prod(out[0].shape)={out_size}, but should be {len(indices)}"
+        else:
+            for i in range(out_size):
+                idx = tuple(int(x[i]) for x in out)
+                f_idx = f"Extrapolated index (x[{i}] for x in out)={idx}"
+                f_element = f"x[{idx}]={x[idx]}"
+                assert idx in indices, f"{f_idx} results in {f_element}, a zero element"
+                assert (
+                    idx == indices[i]
+                ), f"{f_idx} is in the wrong position, should be {indices.index(idx)}"
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise
 
 
 @given(
@@ -188,31 +208,36 @@ def test_where(shapes, dtypes, data):
     x1 = data.draw(hh.arrays(dtype=dtypes[0], shape=shapes[1]), label="x1")
     x2 = data.draw(hh.arrays(dtype=dtypes[1], shape=shapes[2]), label="x2")
 
-    out = xp.where(cond, x1, x2)
-
-    shape = sh.broadcast_shapes(*shapes)
-    ph.assert_shape("where", out_shape=out.shape, expected=shape)
-    # TODO: generate indices without broadcasting arrays
-    _cond = xp.broadcast_to(cond, shape)
-    _x1 = xp.broadcast_to(x1, shape)
-    _x2 = xp.broadcast_to(x2, shape)
-    for idx in sh.ndindex(shape):
-        if _cond[idx]:
-            ph.assert_0d_equals(
-                "where",
-                x_repr=f"_x1[{idx}]",
-                x_val=_x1[idx],
-                out_repr=f"out[{idx}]",
-                out_val=out[idx]
-            )
-        else:
-            ph.assert_0d_equals(
-                "where",
-                x_repr=f"_x2[{idx}]",
-                x_val=_x2[idx],
-                out_repr=f"out[{idx}]",
-                out_val=out[idx]
-            )
+    repro_snippet = ph.format_snippet(f"xp.where({cond!r}, {x1!r}, {x2!r})")
+    try:
+        out = xp.where(cond, x1, x2)
+
+        shape = sh.broadcast_shapes(*shapes)
+        ph.assert_shape("where", out_shape=out.shape, expected=shape)
+        # TODO: generate indices without broadcasting arrays
+        _cond = xp.broadcast_to(cond, shape)
+        _x1 = xp.broadcast_to(x1, shape)
+        _x2 = xp.broadcast_to(x2, shape)
+        for idx in sh.ndindex(shape):
+            if _cond[idx]:
+                ph.assert_0d_equals(
+                    "where",
+                    x_repr=f"_x1[{idx}]",
+                    x_val=_x1[idx],
+                    out_repr=f"out[{idx}]",
+                    out_val=out[idx]
+                )
+            else:
+                ph.assert_0d_equals(
+                    "where",
+                    x_repr=f"_x2[{idx}]",
+                    x_val=_x2[idx],
+                    out_repr=f"out[{idx}]",
+                    out_val=out[idx]
+                )
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise
 
 
 @pytest.mark.min_version("2023.12")
@@ -238,12 +263,17 @@ def test_searchsorted(data):
         label="x2",
     )
 
-    out = xp.searchsorted(x1, x2, sorter=sorter)
+    repro_snippet = ph.format_snippet(f"xp.searchsorted({x1!r}, {x2!r}, sorter={sorter!r})")
+    try:
+        out = xp.searchsorted(x1, x2, sorter=sorter)
 
-    ph.assert_dtype(
-        "searchsorted",
-        in_dtype=[x1.dtype, x2.dtype],
-        out_dtype=out.dtype,
-        expected=xp.__array_namespace_info__().default_dtypes()["indexing"],
-    )
-    # TODO: shapes and values testing
+        ph.assert_dtype(
+            "searchsorted",
+            in_dtype=[x1.dtype, x2.dtype],
+            out_dtype=out.dtype,
+            expected=xp.__array_namespace_info__().default_dtypes()["indexing"],
+        )
+        # TODO: shapes and values testing
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise