tril and triu fail on 1×1 matrices

[zackmdavis]

I just upgraded to my project to use ndarray master in order to be able to use #1386, but discovered that these are failing with thread [test name] has overflowed its stack. (It's a two-dimensional array where both dimensions are 1, perhaps an easy edge-case to overlook.)

I was about to say, "I guess I'm glad I caught this before the eventual 0.16 release", but it looks like 0.16 came out four hours ago?

    #[test]
    fn test_tril() {
        let a: Array2<f32> = Array2::ones((1, 1));
        let b = a.tril(0);
    }

    #[test]
    fn test_triu() {
        let a: Array2<f32> = Array2::ones((1, 1));
        let b = a.triu(0);
    }

[akern40]

Well that's definitely on me! I'm try to take a look at it this week. Just out of curiosity, what would you as a user expect as the output of those?

[zackmdavis]

@akern40 The documentation's promise of "the array with elements below/above the kth diagonal zeroed" would seem to imply that the sole entry of array![[1]].triu(k) should be 1 when k is zero or negative and 0 when k is positive, and that of array![[1]].tril(k) should be 1 when k is zero or positive and 0 when k is negative—and that is the behavior I'm seeing from PyTorch.

The context in which this bit me was while implementing attention in a neural network model: there's a part where you're supposed to mask everything above the main diagonal of a square matrix, but if there was only one input token, then the matrix is 1×1.