code: def soft_argmax(heatmaps, joint_num):
      dims = heatmaps.shape
      depth_dim = (int)(dims[1] / joint_num)
      heatmaps = heatmaps.reshape(
          (-1, joint_num, depth_dim * dims[2] * dims[3])
      )
      heatmaps = F.softmax(heatmaps, 2)
      heatmaps = heatmaps.reshape(
          (-1, joint_num, depth_dim, dims[2], dims[3])
      )
  
      accu_x = heatmaps.sum(axis=(2, 3))
      accu_y = heatmaps.sum(axis=(2, 4))
      accu_z = heatmaps.sum(axis=(3, 4))
  
      accu_x = accu_x * paddle.arange(1, 33)
      accu_y = accu_y * paddle.arange(1, 33)
      accu_z = accu_z * paddle.arange(1, 33)
  
      accu_x = accu_x.sum(axis=2, keepdim=True) - 1
      accu_y = accu_y.sum(axis=2, keepdim=True) - 1
      accu_z = accu_z.sum(axis=2, keepdim=True) - 1
  
      coord_out = paddle.concat(
          (accu_x, accu_y, accu_z), axis=2
      )  # [batch_size, joint_num, 3]
  
      return coord_out
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import paddle
  import numpy
  import math
  import cv2
  from ppdet.core.workspace import register, create
  from .meta_arch import BaseArch
  from ..keypoint_utils import transform_preds
  from ..None import layers
  from paddle.nn import functional
label_desc: Use image node_id 1 for calling a global function with example usage: soft_argmax(heatmaps, joint_num) and returns: coord_out