Question on Soft argmax implementations.

[widiba03304]

AFAIK, soft argmax is an expectation of heatmap, which should be calculated without control flows like if or for. However, the code contains some control flows for tensorflow and pytorch impl.

• https://github.com/isarandi/metrabs/blob/8b2b116dd27372e7dbd8207809f868df4e3f852e/metrabs_tf/tfu.py#L220
• https://github.com/isarandi/metrabs/blob/8b2b116dd27372e7dbd8207809f868df4e3f852e/metrabs_pytorch/ptu.py#L58

Can you check if my understanding is correct?

1. Let's assume an input with a shape [B, D, J, H, W].
2. The softmax to the activation reduces the axes [D, H, W], which results in the sum of [D, H, W] is 1.
3. However, these implementations multiply indices (linspace) for [D, H], [D, W], and [H, W], not [D, H, W].

Why is the softmax performed for [D, H, W] if the indices for [D, H], [D, W], [H, W] should be multiplied by it?

Thanks in advance.