Proposal: add support for returning upper and lower triangles of an array

[kgryte]

Proposal

Add APIs for returning the upper (triu) and lower (tril) triangles of an array.

Motivation

Currently, the array API specification does not provide APIs for returning the upper and lower triangles of an array.

These APIs were identified as potential candidates in gh-187 for inclusion in the next version of the standard. However, depending on the outcome of gh-217 may be necessary to include in the current version in order to ensure users can maintain equivalent functionality.

Background

triu and tril APIs are common across array libraries (see here). The notable exception being TensorFlow. MXNet seems to have APIs for returning the triangles of an array, but these are not documented in public documentation, so were not picked up during API comparison analysis.

Both triu and tril should be straightforward to include in the specification, as the signatures are consistent across libraries.

tril(m, k=0) → array
triu(m, k=0) → array

where k specifies the diagonal above/below which to zero elements.

[leofang]

@kgryte You're 10 mins ahead of me! I was working on cupy/cupy#5555 and just realized the very same need. I think we need them to be able to handle and return a batch of matrices like all other linalg functions, which is currently not possible with NumPy/CuPy.

[kgryte]

@leofang I think handling batches could make sense; however, I don't think any array libraries currently support batching, including Torch. Should be a straightforward extension, as should not conflict with any existing NumPy et al usage patterns.

[lezcano]

The docs of tril and triu are outdated, PyTorch does support batched inputs for these functions. I agree with @leofang that it would be good to write the specification for batched inputs.

[kgryte]

I've opened gh-243 which adds both tril and triu to the specification, with support for providing a stack of matrices.