# Averaging non-zero tensors along an axis

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] ]
]
``````

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

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]) )
``````

Oh right, second axis makes sense.

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

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