def _get_batch_topk(self, batch_heatmaps: Tensor, batch_tags: Tensor,
                        k: int):
        """Get top-k response values from the heatmaps and corresponding tag
        values from the tagging heatmaps.

        Args:
            batch_heatmaps (Tensor): Keypoint detection heatmaps in shape
                (B, K, H, W)
            batch_tags (Tensor): Tagging heatmaps in shape (B, C, H, W), where
                the tag dim C is 2*K when using flip testing, or K otherwise
            k (int): The number of top responses to get

        Returns:
            tuple:
            - topk_vals (Tensor): Top-k response values of each heatmap in
                shape (B, K, Topk)
            - topk_tags (Tensor): The corresponding embedding tags of the
                top-k responses, in shape (B, K, Topk, L)
            - topk_locs (Tensor): The location of the top-k responses in each
                heatmap, in shape (B, K, Topk, 2) where last dimension
                represents x and y coordinates
        """
        B, K, H, W = batch_heatmaps.shape
        L = batch_tags.shape[1] // K

        # shape of topk_val, top_indices: (B, K, TopK)
        topk_vals, topk_indices = batch_heatmaps.flatten(-2, -1).topk(
            k, dim=-1)

        topk_tags_per_kpts = [
            torch.gather(_tag, dim=2, index=topk_indices)
            for _tag in torch.unbind(batch_tags.view(B, L, K, H * W), dim=1)
        ]

        topk_tags = torch.stack(topk_tags_per_kpts, dim=-1)  # (B, K, TopK, L)
        topk_locs = torch.stack([topk_indices % W, topk_indices // W],
                                dim=-1)  # (B, K, TopK, 2)

        return topk_vals, topk_tags, topk_locs
Get top-k response values from the heatmaps and corresponding tag
        values from the tagging heatmaps.

        Args:
            batch_heatmaps (Tensor): Keypoint detection heatmaps in shape
                (B, K, H, W)
            batch_tags (Tensor): Tagging heatmaps in shape (B, C, H, W), where
                the tag dim C is 2*K when using flip testing, or K otherwise
            k (int): The number of top responses to get

        Returns:
            tuple:
            - topk_vals (Tensor): Top-k response values of each heatmap in
                shape (B, K, Topk)
            - topk_tags (Tensor): The corresponding embedding tags of the
                top-k responses, in shape (B, K, Topk, L)
            - topk_locs (Tensor): The location of the top-k responses in each
                heatmap, in shape (B, K, Topk, 2) where last dimension
                represents x and y coordinates
        
mmpose/mmpose/codecs/associative_embedding.py