From d888d62548ddb1354bc61cd8e01567a70d49987c Mon Sep 17 00:00:00 2001 From: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Date: Mon, 13 Jul 2026 03:46:33 +0530 Subject: [PATCH 1/2] Adding tests for tril, triu and adding dtype parameter for linalg.trace Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> --- array_api_tests/test_linalg.py | 10 +++-- .../test_manipulation_functions.py | 41 +++++++++++++++++++ 2 files changed, 47 insertions(+), 4 deletions(-) diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py index e73bc0a9..29d7775b 100644 --- a/array_api_tests/test_linalg.py +++ b/array_api_tests/test_linalg.py @@ -920,7 +920,8 @@ def test_tensordot(x1, x2, kw): x=arrays(dtype=numeric_dtypes, shape=matrix_shapes()), # offset may produce an overflow if it is too large. Supporting offsets # that are way larger than the array shape isn't very important. - kw=kwargs(offset=integers(-MAX_ARRAY_SIZE, MAX_ARRAY_SIZE)) + kw=kwargs(offset=integers(-MAX_ARRAY_SIZE, MAX_ARRAY_SIZE), + dtype=sampled_from(dh.numeric_dtypes)) ) def test_trace(x, kw): res = linalg.trace(x, **kw) @@ -939,10 +940,11 @@ def test_trace(x, kw): n, m = x.shape[-2:] ph.assert_result_shape('trace', x.shape, res.shape, expected=x.shape[:-2]) - def true_trace(x_stack, offset=0): + def true_trace(x_stack, offset=0, dtype=None): # Note: the spec does not specify that offset must be within the # bounds of the matrix. A large offset should just produce a size 0 # diagonal in the last dimension (trace 0). See test_diagonal(). + out_dtype = x.dtype if dtype is None else dtype if offset < 0: diag_size = min(n, m, max(n + offset, 0)) elif offset == 0: @@ -954,8 +956,8 @@ def true_trace(x_stack, offset=0): x_stack_diag = [x_stack[i, i + offset] for i in range(diag_size)] else: x_stack_diag = [x_stack[i - offset, i] for i in range(diag_size)] - result = xp.asarray(xp.stack(x_stack_diag) if x_stack_diag else [], dtype=x.dtype) - return _array_module.sum(result) + result = xp.asarray(xp.stack(x_stack_diag) if x_stack_diag else [], dtype=out_dtype) + return _array_module.sum(result, dtype=dtype) _test_stacks(linalg.trace, x, **kw, res=res, dims=0, true_val=true_trace) diff --git a/array_api_tests/test_manipulation_functions.py b/array_api_tests/test_manipulation_functions.py index e886629b..1b03a3e1 100644 --- a/array_api_tests/test_manipulation_functions.py +++ b/array_api_tests/test_manipulation_functions.py @@ -613,3 +613,44 @@ def test_unstack(x, data): except Exception as exc: ph.add_note(exc, repro_snippet) 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 From ce057d38b311f9a295cf5fb3d897ab1b22cbe005 Mon Sep 17 00:00:00 2001 From: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Date: Mon, 13 Jul 2026 17:33:27 +0530 Subject: [PATCH 2/2] Moving tests for tril and triu to another PR Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> --- .../test_manipulation_functions.py | 41 ------------------- 1 file changed, 41 deletions(-) diff --git a/array_api_tests/test_manipulation_functions.py b/array_api_tests/test_manipulation_functions.py index 1b03a3e1..e886629b 100644 --- a/array_api_tests/test_manipulation_functions.py +++ b/array_api_tests/test_manipulation_functions.py @@ -613,44 +613,3 @@ def test_unstack(x, data): except Exception as exc: ph.add_note(exc, repro_snippet) 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