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