Hi everyone,
I try to implement the following function:

image.jpg1014×196 24.5 KB

At this stage, I have e.g. a tensor [[1,0,3], [0, 1, 2], [3, 2, 1]] and the softmax should be done for only values greater than 0 (neighbors), thus [[0.12, 0, 0.88], [0, 0.27, 0.73], [0.66, 0.24, 0.1]].
Is there an existing method to do this or does anyone have an idea ?
Thank you for your help !