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.

• metrabs/metrabs_tf/tfu.py

  Line 220 in 8b2b116

  def decode_heatmap(inp, axis=-1, output_coord_axis=-1):

• metrabs/metrabs_pytorch/ptu.py

  Line 58 in 8b2b116

  def decode_heatmap(inp, dim=-1, output_coord_dim=-1):

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.