Averaging non-zero tensors along an axis

1

Hello everybody!

I am wondering how we can average on the second dimension of a tensor without taking zero arrays into account. See below:

A = (tensor) [ [ [1, 2, 3], [0, 0 , 0], [3, 4, 5] ],
               [ [0, 0, 0], [5, 6, 7], [0, 1, 2] ]  
              ]

Avg(A) = (tensor) [  [ [2, 3, 4] ], 
                     [ [2.5, 3.5, 4.5] ]  
                  ]
2

Hi @sajastu,

Can you explain how do you obtain 3 for Avg(A[0][1]) when A[0][1] == [0, 0, 0]?
I would expect this output:

Avg(A) = [[2, 0, 4  ],
          [0, 6, 1.5]]

There might be (I hope) a cleaner way, but you can do this:

avg = A.sum(axis=2) / (A != 0).sum(axis=2).clamp(min=1).float()

(clamp to avoid division by zero)

EDIT: whoops, I misread “without taking zero in arrays into account”, above solution does not count. Still I am curious to know how you get this output

3

Thanks for the reply @spanev

I should edit the question, I want to calculate the mean along the second axis. The first array (first row) in tensor Avg(A) is calculated by averaging two non-zero arrays in tensor A. As such:

[2, 3, 4]  =  Avg( ([1, 2, 3], [3, 4, 5]) )
[2.5, 3.5, 4.5]  =  Avg( ([5, 6, 7], [0, 1, 2]) )
4

Oh right, second axis makes sense. :smile:

Well, this should work:

avg = A.sum(axis=1) / (A.sum(axis=2) > 0).sum(axis=1).view(-1,1).float()

In the denominator, we are counting the number of non-zero arrays on the second axis. We can then rely on broadcasting for the division.

2 Likes
5

To get the non zero indices you could use .nonzero(as_tuple=True)

x[x.nonzero(as_tuple=True)].view(-1, x.shape[1])

will give you back the original tensor with all the zero tensors removed. Then you could just use .mean(dim=1) as usual