6

Given input like:

tensor([[[1.9392, -1.9266,  0.9664],
         [0.0000, -1.9266,  0.9664],
         [0.0000, -0.0000,  0.9664]]])

My desired output is:

tensor([[[0.4596,  0.0096, 0.1737],
         [0.0000,  0.0096, 0.1737],
         [0.0000, -0.0000, 0.1737]]])

I.e. just calculating the function over the upper triangular elements.

iacob
  • 20,084
  • 6
  • 92
  • 119
shu.liu
  • 63
  • 3

1 Answers1

5

You can access the upper triangular elements with torch.triu_indices:

t = tensor([[1.9392, -1.9266, 0.9664], 
            [0.0000, -1.9266, 0.9664], 
            [0.0000, -0.0000, 0.9664]]) 

idx = torch.triu_indices(*t.shape)
soft = F.softmax(t[idx[0], idx[1]], dim=0)

If you want to reassign the values as in your desired output:

>>> t[idx[0], idx[1]] = soft
>>> t

tensor([[0.4596,  0.0096, 0.1737],
        [0.0000,  0.0096, 0.1737],
        [0.0000, -0.0000, 0.1737]])
iacob
  • 20,084
  • 6
  • 92
  • 119
  • if tensor shape is 3D, for example r.shape=torch.tensor([5,3,3]) i think use for loop like this `for i in range(r.shape[0]): sub_r = r[i] other code like you` it's right? – shu.liu Apr 20 '21 at 01:43