def predict_all(self,
                    feats: Dict[str, torch.Tensor],
                    data_samples: List[DataSample],
                    num_images: int = None,
                    num_texts: int = None,
                    cal_i2t: bool = True,
                    cal_t2i: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute similarity matrix between images and texts across all ranks.

        Args:
            feats (Dict[str, torch.Tensor]): Features from the current rank.
            data_samples (List[DataSample]): Data samples from the current
                rank.
            num_images (int, optional): Number of images to use.
                Defaults to None.
            num_texts (int, optional): Number of texts to use.
                Defaults to None.
            cal_i2t (bool, optional): Whether to compute image-to-text
                similarity. Defaults to True.
            cal_t2i (bool, optional): Whether to compute text-to-image
                similarity. Defaults to True.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: Image-to-text and text-to-image
            similarity matrices.
        """
        text_ids = feats['text_ids']
        text_attn_mask = feats['text_attn_mask']
        image_embeds = feats.get('image_embeds', None)
        image_feat = feats['image_feat']
        text_feat = feats['text_feat']

        num_images = num_images or image_feat.size(0)
        num_texts = num_texts or text_feat.size(0)

        if not self.fast_match:
            image_embeds_all = all_gather_concat(image_embeds)[:num_images]
        else:
            image_embeds_all = None
        image_feat_all = all_gather_concat(image_feat)[:num_images]
        text_feat_all = all_gather_concat(text_feat)[:num_texts]
        text_ids_all = all_gather_concat(text_ids)[:num_texts]
        text_attn_mask_all = all_gather_concat(text_attn_mask)[:num_texts]

        results = []
        if cal_i2t:
            result_i2t = self.compute_score_matrix_i2t(
                image_feat,
                image_embeds,
                text_feat_all,
                text_ids_all,
                text_attn_mask_all,
            )
            results.append(
                self._get_predictions(result_i2t, data_samples, mode='i2t'))
        if cal_t2i:
            result_t2i = self.compute_score_matrix_t2i(
                image_feat_all,
                image_embeds_all,
                text_feat,
                text_ids,
                text_attn_mask,
            )
            results.append(
                self._get_predictions(result_t2i, data_samples, mode='t2i'))
        return tuple(results)
Compute similarity matrix between images and texts across all ranks.

        Args:
            feats (Dict[str, torch.Tensor]): Features from the current rank.
            data_samples (List[DataSample]): Data samples from the current
                rank.
            num_images (int, optional): Number of images to use.
                Defaults to None.
            num_texts (int, optional): Number of texts to use.
                Defaults to None.
            cal_i2t (bool, optional): Whether to compute image-to-text
                similarity. Defaults to True.
            cal_t2i (bool, optional): Whether to compute text-to-image
                similarity. Defaults to True.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: Image-to-text and text-to-image
            similarity matrices.
        
mmpretrain/models/multimodal/blip2/blip2_retriever.py