Skip to content

ENH: Adding tests for tril and triu #442

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
prady0t wants to merge 2 commits into
base: master
Choose a base branch
from
+51 −0
Draft

ENH: Adding tests for tril and triu #442

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5ca8a3b
Adding tests for tril and triu
prady0t Jul 13, 2026
409ac69
remove in-place modifications to address jax.numpy
prady0t Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions array_api_tests/test_creation_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -392,6 +392,57 @@ def test_eye(n_rows, n_cols, kw):
raise



@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)
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,
)
Comment on lines +408 to +415

@prady0t prady0t Jul 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried to test the tril and triu by essentially looping over idx, so if it has a lot of dimensions, the test becomes slow, and complexity increases subsequently. As an example, running with --max-examples=10_000 tril and triu takes over 5 minutes!

Note: This is valid for jax.numpy. Testing with --max-examples=10_000 for array_api_compat.numpy takes 30 seconds.

@prady0t prady0t Jul 16, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tagging @ev-br

@ev-br ev-br Jul 16, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This calls ph.assert_array_elements in the loop. Could we lift it out of the loop to make it closer to

https://github.com/data-apis/array-api-tests/blob/master/array_api_tests/test_creation_functions.py#L383-L389

There'll still be a loop to construct expected, yes, and it's not necessarily 2D, so it'll take some gymnastics to assemble approriately nested list. This exercise will tell us if the runtime is sensible or if we need to tune it further (thread max_dim or max_size through matrix_shapes, apply @univectorized to decimate the number of examples or something else---no point messing with either of these if we can assemble expected).

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)
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


default_unsafe_dtypes = [xp.uint64]
if dh.default_int == xp.int32:
default_unsafe_dtypes.extend([xp.uint32, xp.int64])
Expand Down
Loading