tf.matrix_band_part equivalent

1

I need to create upper triangular masking tensor in pytorch. in tensorflow it’s easy using matrix_band_part. is it any pytorch equivalent for this function?
it’s like numpy triu_indices but I need for tensor, not just matrix.

1 Like
2

Hi,

Isn’t that what torch.triu() is doing?

3

thank you for answer but torch.triu() is also only works with matrix.

5

Did you find any solutions for this?

6

no. I did a for loop for that task at that time. I dont know if they have added any new possibility or not.

1 Like
7

The batch version of .triu() has been added since then :slight_smile:

8

I just tried it yesterday with PyTorch 1.0.0, is that for a newer version? If not, I’ll leave a code to reproduce my error here.

9

We’re at 1.3.1 now :slight_smile: I just tried again with latest version and it works:

>>> a = torch.rand(2, 4, 4)
>>> a.triu()
tensor([[[0.6312, 0.0578, 0.9598, 0.4663],
         [0.0000, 0.3202, 0.9285, 0.6618],
         [0.0000, 0.0000, 0.4512, 0.0584],
         [0.0000, 0.0000, 0.0000, 0.2116]],

        [[0.1818, 0.8329, 0.9879, 0.6801],
         [0.0000, 0.2870, 0.3073, 0.3876],
         [0.0000, 0.0000, 0.6292, 0.7130],
         [0.0000, 0.0000, 0.0000, 0.7048]]])
1 Like
10

Just upgraded and works as expected. Thanks for the tip!

1 Like