ENH: add "repro_snippets" to test_indexing_functions.py

Author: ev-br

array_api_tests/test_indexing_functions.py

@@ -35,39 +35,43 @@ def test_take(x, data):
     indices = xp.asarray(_indices, dtype=dh.default_int)
     note(f"{indices=}")
 
-    out = xp.take(x, indices, **kw)
-
-    ph.assert_dtype("take", in_dtype=x.dtype, out_dtype=out.dtype)
-    ph.assert_shape(
-        "take",
-        out_shape=out.shape,
-        expected=x.shape[:n_axis] + (len(_indices),) + x.shape[n_axis + 1:],
-        kw=dict(
-            x=x,
-            indices=indices,
-            axis=axis,
-        ),
-    )
-    out_indices = sh.ndindex(out.shape)
-    axis_indices = list(sh.axis_ndindex(x.shape, n_axis))
-    for axis_idx in axis_indices:
-        f_axis_idx = sh.fmt_idx("x", axis_idx)
-        for i in _indices:
-            f_take_idx = sh.fmt_idx(f_axis_idx, i)
-            indexed_x = x[axis_idx][i, ...]
-            for at_idx in sh.ndindex(indexed_x.shape):
-                out_idx = next(out_indices)
-                ph.assert_0d_equals(
-                    "take",
-                    x_repr=sh.fmt_idx(f_take_idx, at_idx),
-                    x_val=indexed_x[at_idx],
-                    out_repr=sh.fmt_idx("out", out_idx),
-                    out_val=out[out_idx],
-                )
-    # sanity check
-    with pytest.raises(StopIteration):
-        next(out_indices)
+    repro_snippet = ph.format_snippet(f"xp.take({x!r}, {indices!r}, **kw) with {kw = }")
+    try:
+        out = xp.take(x, indices, **kw)
 
+        ph.assert_dtype("take", in_dtype=x.dtype, out_dtype=out.dtype)
+        ph.assert_shape(
+            "take",
+            out_shape=out.shape,
+            expected=x.shape[:n_axis] + (len(_indices),) + x.shape[n_axis + 1:],
+            kw=dict(
+                x=x,
+                indices=indices,
+                axis=axis,
+            ),
+        )
+        out_indices = sh.ndindex(out.shape)
+        axis_indices = list(sh.axis_ndindex(x.shape, n_axis))
+        for axis_idx in axis_indices:
+            f_axis_idx = sh.fmt_idx("x", axis_idx)
+            for i in _indices:
+                f_take_idx = sh.fmt_idx(f_axis_idx, i)
+                indexed_x = x[axis_idx][i, ...]
+                for at_idx in sh.ndindex(indexed_x.shape):
+                    out_idx = next(out_indices)
+                    ph.assert_0d_equals(
+                        "take",
+                        x_repr=sh.fmt_idx(f_take_idx, at_idx),
+                        x_val=indexed_x[at_idx],
+                        out_repr=sh.fmt_idx("out", out_idx),
+                        out_val=out[out_idx],
+                    )
+        # sanity check
+        with pytest.raises(StopIteration):
+            next(out_indices)
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise
 
 @pytest.mark.unvectorized
 @pytest.mark.min_version("2024.12")
@@ -103,26 +107,33 @@ def test_take_along_axis(x, data):
     )
     note(f"{indices=}  {idx_shape=}")
 
-    out = xp.take_along_axis(x, indices, **axis_kw)
-
-    ph.assert_dtype("take_along_axis", in_dtype=x.dtype, out_dtype=out.dtype)
-    ph.assert_shape(
-        "take_along_axis",
-        out_shape=out.shape,
-        expected=x.shape[:n_axis] + (new_len,) + x.shape[n_axis+1:],
-        kw=dict(
-            x=x,
-            indices=indices,
-            axis=axis,
-        ),
+    repro_snippet = ph.format_snippet(
+        f"xp.take_along_axis({x!r}, {indices!r}, **axis_kw) with {axis_kw = }"
     )
+    try:
+        out = xp.take_along_axis(x, indices, **axis_kw)
+
+        ph.assert_dtype("take_along_axis", in_dtype=x.dtype, out_dtype=out.dtype)
+        ph.assert_shape(
+            "take_along_axis",
+            out_shape=out.shape,
+            expected=x.shape[:n_axis] + (new_len,) + x.shape[n_axis+1:],
+            kw=dict(
+                x=x,
+                indices=indices,
+                axis=axis,
+            ),
+        )
 
-    # value test: notation is from `np.take_along_axis` docstring
-    Ni, Nk = x.shape[:n_axis], x.shape[n_axis+1:]
-    for ii in sh.ndindex(Ni):
-        for kk in sh.ndindex(Nk):
-            a_1d = x[ii + (slice(None),) + kk]
-            i_1d = indices[ii + (slice(None),) + kk]
-            o_1d = out[ii + (slice(None),) + kk]
-            for j in range(new_len):
-                assert o_1d[j] == a_1d[i_1d[j]], f'{ii=}, {kk=}, {j=}'
+        # value test: notation is from `np.take_along_axis` docstring
+        Ni, Nk = x.shape[:n_axis], x.shape[n_axis+1:]
+        for ii in sh.ndindex(Ni):
+            for kk in sh.ndindex(Nk):
+                a_1d = x[ii + (slice(None),) + kk]
+                i_1d = indices[ii + (slice(None),) + kk]
+                o_1d = out[ii + (slice(None),) + kk]
+                for j in range(new_len):
+                    assert o_1d[j] == a_1d[i_1d[j]], f'{ii=}, {kk=}, {j=}'
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise