From 5ca8a3bebaf42b8c2c0390ad08253c0b0d8bcbf6 Mon Sep 17 00:00:00 2001 From: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Date: Mon, 13 Jul 2026 17:29:10 +0530 Subject: [PATCH 1/5] Adding tests for tril and triu Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> --- array_api_tests/test_creation_functions.py | 44 ++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/array_api_tests/test_creation_functions.py b/array_api_tests/test_creation_functions.py index 77e5dbb3..3782c122 100644 --- a/array_api_tests/test_creation_functions.py +++ b/array_api_tests/test_creation_functions.py @@ -392,6 +392,50 @@ def test_eye(n_rows, n_cols, kw): raise +def _triangular_expected(x, *, k, upper): + expected = xp.zeros_like(x) + n, m = x.shape[-2:] + for idx in sh.ndindex(x.shape[:-2]): + for i in range(n): + for j in range(m): + keep = j <= i + k if not upper else j >= i + k + if keep: + expected[idx + (i, j)] = x[idx + (i, j)] + return expected + + +@given(x=hh.arrays(dtype=hh.numeric_dtypes, shape=hh.matrix_shapes()), data=st.data()) +def test_tril(x, data): + n, m = x.shape[-2:] + k = data.draw(st.integers(min_value=-max(n, m, 1), max_value=max(n, m, 1)), label="k") + repro_snippet = ph.format_snippet(f"xp.tril({x!r}, k={k!r})") + try: + out = xp.tril(x, k=k) + ph.assert_dtype("tril", in_dtype=x.dtype, out_dtype=out.dtype) + ph.assert_shape("tril", out_shape=out.shape, expected=x.shape) + expected = _triangular_expected(x, k=k, upper=False) + ph.assert_array_elements("tril", out=out, expected=expected, kw={"k": k}) + except Exception as exc: + ph.add_note(exc, repro_snippet) + raise + + +@given(x=hh.arrays(dtype=hh.numeric_dtypes, shape=hh.matrix_shapes()), data=st.data()) +def test_triu(x, data): + n, m = x.shape[-2:] + k = data.draw(st.integers(min_value=-max(n, m, 1), max_value=max(n, m, 1)), label="k") + repro_snippet = ph.format_snippet(f"xp.triu({x!r}, k={k!r})") + try: + out = xp.triu(x, k=k) + ph.assert_dtype("triu", in_dtype=x.dtype, out_dtype=out.dtype) + ph.assert_shape("triu", out_shape=out.shape, expected=x.shape) + expected = _triangular_expected(x, k=k, upper=True) + ph.assert_array_elements("triu", out=out, expected=expected, kw={"k": k}) + except Exception as exc: + ph.add_note(exc, repro_snippet) + raise + + default_unsafe_dtypes = [xp.uint64] if dh.default_int == xp.int32: default_unsafe_dtypes.extend([xp.uint32, xp.int64]) From 409ac69c46f7d7eea07624b5aea42bfc5e040e3d Mon Sep 17 00:00:00 2001 From: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Date: Thu, 16 Jul 2026 19:00:08 +0530 Subject: [PATCH 2/5] remove in-place modifications to address jax.numpy Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> --- array_api_tests/test_creation_functions.py | 41 +++++++++++++--------- 1 file changed, 24 insertions(+), 17 deletions(-) diff --git a/array_api_tests/test_creation_functions.py b/array_api_tests/test_creation_functions.py index 3782c122..88218291 100644 --- a/array_api_tests/test_creation_functions.py +++ b/array_api_tests/test_creation_functions.py @@ -392,29 +392,27 @@ def test_eye(n_rows, n_cols, kw): raise -def _triangular_expected(x, *, k, upper): - expected = xp.zeros_like(x) - n, m = x.shape[-2:] - for idx in sh.ndindex(x.shape[:-2]): - for i in range(n): - for j in range(m): - keep = j <= i + k if not upper else j >= i + k - if keep: - expected[idx + (i, j)] = x[idx + (i, j)] - return expected - @given(x=hh.arrays(dtype=hh.numeric_dtypes, shape=hh.matrix_shapes()), data=st.data()) def test_tril(x, data): n, m = x.shape[-2:] - k = data.draw(st.integers(min_value=-max(n, m, 1), max_value=max(n, m, 1)), label="k") + k = data.draw( + st.integers(min_value=-max(n, m, 1), max_value=max(n, m, 1)), + label="k", + ) repro_snippet = ph.format_snippet(f"xp.tril({x!r}, k={k!r})") try: out = xp.tril(x, k=k) ph.assert_dtype("tril", in_dtype=x.dtype, out_dtype=out.dtype) ph.assert_shape("tril", out_shape=out.shape, expected=x.shape) - expected = _triangular_expected(x, k=k, upper=False) - ph.assert_array_elements("tril", out=out, expected=expected, kw={"k": k}) + for idx in sh.ndindex(out.shape): + *_, i, j = idx + expected = x[idx] if j <= i + k else xp.asarray(0, dtype=out.dtype) + ph.assert_array_elements( + "tril", + out=out[idx], + expected=expected, + ) except Exception as exc: ph.add_note(exc, repro_snippet) raise @@ -423,14 +421,23 @@ def test_tril(x, data): @given(x=hh.arrays(dtype=hh.numeric_dtypes, shape=hh.matrix_shapes()), data=st.data()) def test_triu(x, data): n, m = x.shape[-2:] - k = data.draw(st.integers(min_value=-max(n, m, 1), max_value=max(n, m, 1)), label="k") + k = data.draw( + st.integers(min_value=-max(n, m, 1), max_value=max(n, m, 1)), + label="k", + ) repro_snippet = ph.format_snippet(f"xp.triu({x!r}, k={k!r})") try: out = xp.triu(x, k=k) ph.assert_dtype("triu", in_dtype=x.dtype, out_dtype=out.dtype) ph.assert_shape("triu", out_shape=out.shape, expected=x.shape) - expected = _triangular_expected(x, k=k, upper=True) - ph.assert_array_elements("triu", out=out, expected=expected, kw={"k": k}) + for idx in sh.ndindex(out.shape): + *_, i, j = idx + expected = x[idx] if j >= i + k else xp.asarray(0, dtype=out.dtype) + ph.assert_array_elements( + "triu", + out=out[idx], + expected=expected, + ) except Exception as exc: ph.add_note(exc, repro_snippet) raise From 64e6b97b99ed3e2863d44b5377aba8f649f5ac4c Mon Sep 17 00:00:00 2001 From: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Date: Fri, 17 Jul 2026 17:59:07 +0530 Subject: [PATCH 3/5] use idx for loop entirely Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> --- array_api_tests/test_creation_functions.py | 44 +++++++++++++--------- 1 file changed, 27 insertions(+), 17 deletions(-) diff --git a/array_api_tests/test_creation_functions.py b/array_api_tests/test_creation_functions.py index 88218291..89b49b9c 100644 --- a/array_api_tests/test_creation_functions.py +++ b/array_api_tests/test_creation_functions.py @@ -255,6 +255,7 @@ def test_asarray_scalars(shape, data): ph.assert_kw_dtype("asarray", kw_dtype=_dtype, out_dtype=out.dtype) ph.assert_shape("asarray", out_shape=out.shape, expected=shape) for idx, v_expect in zip(sh.ndindex(out.shape), _obj): + print(f"out.shape is : {out.shape} idx is: {idx}") v = scalar_type(out[idx]) ph.assert_scalar_equals("asarray", type_=scalar_type, idx=idx, out=v, expected=v_expect, kw=kw) except Exception as exc: @@ -405,14 +406,19 @@ def test_tril(x, data): out = xp.tril(x, k=k) ph.assert_dtype("tril", in_dtype=x.dtype, out_dtype=out.dtype) ph.assert_shape("tril", out_shape=out.shape, expected=x.shape) - for idx in sh.ndindex(out.shape): - *_, i, j = idx - expected = x[idx] if j <= i + k else xp.asarray(0, dtype=out.dtype) - ph.assert_array_elements( - "tril", - out=out[idx], - expected=expected, - ) + expected = xp.asarray( + [ + x[idx] if idx[-1] <= idx[-2] + k else 0 + for idx in sh.ndindex(x.shape) + ], + dtype=out.dtype, + ) + expected = xp.reshape(expected, x.shape) + ph.assert_array_elements( + "tril", + out=out, + expected=expected, + ) except Exception as exc: ph.add_note(exc, repro_snippet) raise @@ -430,19 +436,23 @@ def test_triu(x, data): out = xp.triu(x, k=k) ph.assert_dtype("triu", in_dtype=x.dtype, out_dtype=out.dtype) ph.assert_shape("triu", out_shape=out.shape, expected=x.shape) - for idx in sh.ndindex(out.shape): - *_, i, j = idx - expected = x[idx] if j >= i + k else xp.asarray(0, dtype=out.dtype) - ph.assert_array_elements( - "triu", - out=out[idx], - expected=expected, - ) + expected = xp.asarray( + [ + x[idx] if idx[-1] >= idx[-2] + k else 0 + for idx in sh.ndindex(x.shape) + ], + dtype=out.dtype, + ) + expected = xp.reshape(expected, x.shape) + ph.assert_array_elements( + "triu", + out=out, + expected=expected, + ) except Exception as exc: ph.add_note(exc, repro_snippet) raise - default_unsafe_dtypes = [xp.uint64] if dh.default_int == xp.int32: default_unsafe_dtypes.extend([xp.uint32, xp.int64]) From 84be72a5fdc35f26135c634112f7af9049c780c4 Mon Sep 17 00:00:00 2001 From: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:16:30 +0530 Subject: [PATCH 4/5] use stack for elements Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> --- array_api_tests/test_creation_functions.py | 28 ++++++++++++---------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/array_api_tests/test_creation_functions.py b/array_api_tests/test_creation_functions.py index 89b49b9c..37248b64 100644 --- a/array_api_tests/test_creation_functions.py +++ b/array_api_tests/test_creation_functions.py @@ -406,13 +406,15 @@ def test_tril(x, data): out = xp.tril(x, k=k) ph.assert_dtype("tril", in_dtype=x.dtype, out_dtype=out.dtype) ph.assert_shape("tril", out_shape=out.shape, expected=x.shape) - expected = xp.asarray( - [ - x[idx] if idx[-1] <= idx[-2] + k else 0 + expected_elements = [ + x[idx] if idx[-1] <= idx[-2] + k + else xp.asarray(0, dtype=out.dtype) for idx in sh.ndindex(x.shape) - ], - dtype=out.dtype, - ) + ] + if expected_elements: + expected = xp.stack(expected_elements) + else: + expected = xp.asarray([], dtype=out.dtype) expected = xp.reshape(expected, x.shape) ph.assert_array_elements( "tril", @@ -436,13 +438,15 @@ def test_triu(x, data): out = xp.triu(x, k=k) ph.assert_dtype("triu", in_dtype=x.dtype, out_dtype=out.dtype) ph.assert_shape("triu", out_shape=out.shape, expected=x.shape) - expected = xp.asarray( - [ - x[idx] if idx[-1] >= idx[-2] + k else 0 + expected_elements = [ + x[idx] if idx[-1] >= idx[-2] + k + else xp.asarray(0, dtype=out.dtype) for idx in sh.ndindex(x.shape) - ], - dtype=out.dtype, - ) + ] + if expected_elements: + expected = xp.stack(expected_elements) + else: + expected = xp.asarray([], dtype=out.dtype) expected = xp.reshape(expected, x.shape) ph.assert_array_elements( "triu", From 5dd3b5ed253d8f68aed906fae2b9a8e388ea25c4 Mon Sep 17 00:00:00 2001 From: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Date: Fri, 17 Jul 2026 19:40:37 +0530 Subject: [PATCH 5/5] adding comment + removing debug print + moving zero array out of loop Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> --- array_api_tests/test_creation_functions.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/array_api_tests/test_creation_functions.py b/array_api_tests/test_creation_functions.py index 37248b64..7155d45a 100644 --- a/array_api_tests/test_creation_functions.py +++ b/array_api_tests/test_creation_functions.py @@ -255,7 +255,6 @@ def test_asarray_scalars(shape, data): ph.assert_kw_dtype("asarray", kw_dtype=_dtype, out_dtype=out.dtype) ph.assert_shape("asarray", out_shape=out.shape, expected=shape) for idx, v_expect in zip(sh.ndindex(out.shape), _obj): - print(f"out.shape is : {out.shape} idx is: {idx}") v = scalar_type(out[idx]) ph.assert_scalar_equals("asarray", type_=scalar_type, idx=idx, out=v, expected=v_expect, kw=kw) except Exception as exc: @@ -406,14 +405,17 @@ def test_tril(x, data): out = xp.tril(x, k=k) ph.assert_dtype("tril", in_dtype=x.dtype, out_dtype=out.dtype) ph.assert_shape("tril", out_shape=out.shape, expected=x.shape) + zero = xp.asarray(0, dtype=out.dtype) expected_elements = [ x[idx] if idx[-1] <= idx[-2] + k - else xp.asarray(0, dtype=out.dtype) + else zero for idx in sh.ndindex(x.shape) ] if expected_elements: expected = xp.stack(expected_elements) else: + # xp.stack does not accept an empty sequence, and the dtype cannot + # be inferred when there are no elements, so use out.dtype explicitly. expected = xp.asarray([], dtype=out.dtype) expected = xp.reshape(expected, x.shape) ph.assert_array_elements( @@ -438,9 +440,10 @@ def test_triu(x, data): out = xp.triu(x, k=k) ph.assert_dtype("triu", in_dtype=x.dtype, out_dtype=out.dtype) ph.assert_shape("triu", out_shape=out.shape, expected=x.shape) + zero = xp.asarray(0, dtype=out.dtype) expected_elements = [ x[idx] if idx[-1] >= idx[-2] + k - else xp.asarray(0, dtype=out.dtype) + else zero for idx in sh.ndindex(x.shape) ] if expected_elements: