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

Merged
ev-br merged 5 commits into from
Jul 17, 2026
+68 −0
Merged

ENH: Adding tests for tril and triu #442

Changes from all commits
Commits
Show all changes
5 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
64e6b97
use idx for loop entirely
prady0t Jul 17, 2026
84be72a
use stack for elements
prady0t Jul 17, 2026
5dd3b5e
adding comment + removing debug print + moving zero array out of loop
prady0t Jul 17, 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
68 changes: 68 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,74 @@ 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)
zero = xp.asarray(0, dtype=out.dtype)
expected_elements = [
x[idx] if idx[-1] <= idx[-2] + k
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)

@ev-br ev-br Jul 17, 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.

To double-check: this branching is needed even if if expected_elements is an empty list because 1) xp.stack does not accept an empty list, and even if it did, 2) it has no way to guess dtype=out.dtype. Would be nice to add a comment, since it's one of things which goes against numpy finger (and eye) memory.

@prady0t prady0t Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Agreed.

expected = xp.reshape(expected, x.shape)
ph.assert_array_elements(
"tril",
out=out,
expected=expected,
)
Comment on lines +409 to +425

@prady0t prady0t Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I've moved the ph.assert_array_elements outside the loop and used idx entirely for one for loop. expected_elements contains a list of arrays which later gets stacked into one using stack. We then reshape it back to its original shape.

I think this brings us closer to how eye is tested.

@prady0t prady0t Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

One thing that fails right now is running for jax.numpy with --max-examples=10_000 gives:

  + Exception Group Traceback (most recent call last):
  |   File "/Users/pradyotranjan/array-api-tests-1/array_api_tests/test_creation_functions.py", line 398, in test_tril
  |     def test_tril(x, data):
  |                    ^^^
  |   File "/opt/homebrew/Caskroom/miniconda/base/envs/array-api-tests/lib/python3.14/site-packages/hypothesis/core.py", line 2274, in wrapped_test
  |     raise the_error_hypothesis_found
  | hypothesis.errors.FlakyFailure: Hypothesis test_tril(x=Array([[[[ 2.5453735e+18-4.9889314e+16j]],
  | 
  |         [[ 1.5772956e+19+5.1305820e+16j]]],
  | 
  | 
  |        [[[ 2.8746495e+16+2.5450868e+16j]],
  | 
  |         [[ 1.8076227e+19+1.7226468e+19j]]],
  | 
  | 
  |        [[[-5.6867708e+16-4.8846053e+16j]],
  | 
  |         [[ 2.6559498e+16+7.6443090e+18j]]]], dtype=complex64), data=data(...)) produces unreliable results: Falsified on the first call but did not on a subsequent one (1 sub-exception)
  | Falsifying example: test_tril(
  |     x=Array([[[[ 2.5453735e+18-4.9889314e+16j]],
  |     
  |             [[ 1.5772956e+19+5.1305820e+16j]]],
  |     
  |     
  |            [[[ 2.8746495e+16+2.5450868e+16j]],
  |     
  |             [[ 1.8076227e+19+1.7226468e+19j]]],
  |     
  |     
  |            [[[-5.6867708e+16-4.8846053e+16j]],
  |     
  |             [[ 2.6559498e+16+7.6443090e+18j]]]], dtype=complex64),
  |     data=data(...),
  | )
  | Draw 1 (k): 0
  | Unreliable test timings! On an initial run, this test took 1058.55ms, which exceeded the deadline of 800.00ms, but on a subsequent run it took 4.21 ms, which did not. If you expect this sort of variability in your test timings, consider turning deadlines off for this test by setting deadline=None.
  +-+---------------- 1 ----------------
    | Traceback (most recent call last):
    |   File "/opt/homebrew/Caskroom/miniconda/base/envs/array-api-tests/lib/python3.14/site-packages/hypothesis/core.py", line 1212, in _execute_once_for_engine
    |     result = self.execute_once(data)
    |   File "/opt/homebrew/Caskroom/miniconda/base/envs/array-api-tests/lib/python3.14/site-packages/hypothesis/core.py", line 1150, in execute_once
    |     result = self.test_runner(data, run)
    |   File "/opt/homebrew/Caskroom/miniconda/base/envs/array-api-tests/lib/python3.14/site-packages/hypothesis/core.py", line 824, in default_executor
    |     return function(data)
    |   File "/opt/homebrew/Caskroom/miniconda/base/envs/array-api-tests/lib/python3.14/site-packages/hypothesis/core.py", line 1107, in run
    |     return test(*args, **kwargs)
    |   File "/Users/pradyotranjan/array-api-tests-1/array_api_tests/test_creation_functions.py", line 398, in test_tril
    |     def test_tril(x, data):
    |                    ^^^^^^^^
    |   File "/opt/homebrew/Caskroom/miniconda/base/envs/array-api-tests/lib/python3.14/site-packages/hypothesis/core.py", line 1041, in test
    |     raise DeadlineExceeded(
    |         datetime.timedelta(seconds=runtime), self.settings.deadline
    |     )
    | hypothesis.errors.DeadlineExceeded: Test took 1058.55ms, which exceeds the deadline of 800.00ms. If you expect test cases to take this long, you can use @settings(deadline=...) to either set a higher deadline, or to disable it with deadline=None.
    +------------------------------------
===================================================================================================== short test summary info =====================================================================================================
FAILED array_api_tests/test_creation_functions.py::test_tril - DeadlineExceeded('Test took 1058.55ms, which exceeds the deadline of 800.00ms. If you expect test cases to take this long, you can use @settings(deadline=...) to either set a higher deadline, or to disable it with deadline...

@ev-br ev-br Jul 17, 2026 •
edited
Loading

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.

Does it still fail if you add --disable-deadline?

@prady0t prady0t Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Nope. It takes 13 min to complete the test on my machine.

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)
zero = xp.asarray(0, dtype=out.dtype)
expected_elements = [
x[idx] if idx[-1] >= idx[-2] + k
else zero
for idx in sh.ndindex(x.shape)
]
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",
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])
Expand Down
Loading