Make test_eye more efficient

Author: jakevdp

array_api_tests/test_creation_functions.py

@@ -355,13 +355,14 @@ def test_eye(n_rows, n_cols, kw):
     _n_cols = n_rows if n_cols is None else n_cols
     ph.assert_shape("eye", out_shape=out.shape, expected=(n_rows, _n_cols), kw=dict(n_rows=n_rows, n_cols=n_cols))
     f_func = f"[eye({n_rows=}, {n_cols=})]"
-    for i in range(n_rows):
-        for j in range(_n_cols):
+    k = kw.get("k", 0)
+    expected = xp.asarray([[1 if j - i == k else 0
+                            for j in range(_n_cols)] for i in range(n_rows)]).reshape(n_rows, _n_cols)
+    assert out.shape == expected.shape
+    if xp.any(out != expected):
+        for i, j in zip(*xp.where(out != expected)):
             f_indexed_out = f"out[{i}, {j}]={out[i, j]}"
-            if j - i == kw.get("k", 0):
-                assert out[i, j] == 1, f"{f_indexed_out}, should be 1 {f_func}"
-            else:
-                assert out[i, j] == 0, f"{f_indexed_out}, should be 0 {f_func}"
+            assert out[i, j] == expected[i, j], f"{f_indexed_out}, should be {expected[i, j]} {f_func}"
 
 
 default_unsafe_dtypes = [xp.uint64]