    def forward(self, text, experts, ind, raw_captions):
        """Compute joint embeddings and, if requested, a confusion matrix between
        video and text representations in the minibatch.

        Notation: B = batch size, M = number of modalities
        """
        if "retrieval" in self.task:
            # Pass text embeddings through gated units
            text_embd = {}

            # Unroll repeated captions into present minibatch
            bb, captions_per_video, feat_dim = text.size()
            text = text.view(bb * captions_per_video, feat_dim)
            for modality, layer in zip(self.modalities, self.text_GU):
                # NOTE: Due to the batch norm, the gated units are sensitive to passing
                # in a lot of zeroes, so we do the masking step after the forwards pass
                text_ = layer(text)

                # We always assume that text is available for retrieval
                text_ = text_.view(bb, captions_per_video, -1)

                if "text" in self.random_feats:
                    text_ = torch.rand_like(text_)
                text_embd[modality] = text_
            text = text.view(bb, captions_per_video, -1)

            # vladded nans are handled earlier (during pooling)
            # We also avoid zeroing random features, since this will leak information
            # exclude = list(self.vlad_feat_sizes.keys()) + list(self.random_feats)
            # experts = self.mask_missing_embeddings(experts, ind, exclude=exclude)

            # MOE weights computation + normalization - note that we use the first caption
            # sample to predict the weights
            moe_weights = self.compute_moe_weights(text, ind=ind)

        if self.l2renorm:
            for modality in self.modalities:
                norm = experts[modality].norm(p=2, dim=-1, keepdim=True)
                experts[modality] = experts[modality].div(norm)

        for modality, layer in zip(self.modalities, self.trn_list):
            experts[modality] = layer(experts[modality])

        if hasattr(self, "video_dim_reduce"):
            # Embed all features to a common dimension
            for modality, layer in zip(self.modalities, self.video_dim_reduce):
                experts[modality] = layer(experts[modality])

        if self.use_ce:
            dev = experts[self.modalities[0]].device
            if self.include_self:
                all_combinations = list(itertools.product(experts, repeat=2))
            else:
                all_combinations = list(itertools.permutations(experts, 2))
            assert len(self.modalities) > 1, "use_ce requires multiple modalities"

            if self.use_ce in {"pairwise-star", "pairwise-star-specific",
                               "pairwise-star-tensor"}:
                sum_all = 0
                sum_ind = 0
                for mod0 in experts.keys():
                    sum_all += (experts[mod0] * ind[mod0].float().to(dev).unsqueeze(1))
                    sum_ind += ind[mod0].float().to(dev).unsqueeze(1)
                avg_modality = sum_all / sum_ind

            for ii, l in enumerate(self.video_GU):

                mask_num = 0
                curr_mask = 0
                temp_dict = {}
                avai_dict = {}
                curr_modality = self.modalities[ii]

                if self.use_ce == "pairwise-star":
                    fused = torch.cat((experts[curr_modality], avg_modality), 1)  # -> B x 2D
                    temp = self.g_reason_1(fused)  # B x 2D -> B x D
                    temp = self.g_reason_shared(temp)  # B x D -> B x D
                    curr_mask = temp * ind[curr_modality].float().to(dev).unsqueeze(1)

                elif self.use_ce == "pairwise-star-specific":
                    fused = torch.cat((experts[curr_modality], avg_modality), 1)  # -> B x 2D
                    temp = self.g_reason_unshared_weights[ii](fused)
                    temp = self.g_reason_shared(temp)  # B x D -> B x D
                    curr_mask = temp * ind[curr_modality].float().to(dev).unsqueeze(1)

                elif self.use_ce == "pairwise-star-tensor":
                    mod0_reduce = self.dim_reduce(experts[curr_modality])
                    mod0_reduce = mod0_reduce.unsqueeze(2)  # B x reduced_dim x1
                    mod1_reduce = self.dim_reduce(avg_modality)
                    mod1_reduce = mod1_reduce.unsqueeze(1)  # B x1 x reduced_dim
                    flat_dim = self.reduce_dim * self.reduce_dim
                    fused = torch.matmul(mod0_reduce, mod1_reduce).view(-1, flat_dim)
                    temp = self.g_reason_1(fused)  # B x 2D -> B x D
                    temp = self.g_reason_shared(temp)  # B x D -> B x D
                    curr_mask = temp * ind[curr_modality].float().to(dev).unsqueeze(1)

                elif self.use_ce in {"pairwise", "triplet"}:
                    for modality_pair in all_combinations:
                        mod0, mod1 = modality_pair
                        if self.use_ce == "pairwise":
                            if mod0 == curr_modality:
                                new_key = f"{mod0}_{mod1}"
                                fused = torch.cat((experts[mod0], experts[mod1]), 1)
                                temp = self.g_reason_1(fused)  # B x 2D -> B x D
                                temp = self.g_reason_shared(temp)
                                temp_dict[new_key] = temp
                                avail = (ind[mod0].float() * ind[mod1].float())
                                avai_dict[new_key] = avail.to(dev)
                        elif self.use_ce == "triplet":
                            if (curr_modality not in {mod0, mod1}) or self.include_self:
                                new_key = f"{curr_modality}_{mod0}_{mod1}"
                                fused = torch.cat((experts[curr_modality], experts[mod0],
                                                   experts[mod1]), 1)  # -> B x 2D
                                temp = self.g_reason_1(fused)  # B x 2D -> B x D
                                temp = self.g_reason_shared(temp)
                                temp_dict[new_key] = temp
                                avail = (ind[curr_modality].float() * ind[mod0].float() *
                                         ind[mod1].float()).to(dev)
                                avai_dict[new_key] = avail

                    # Combine the paired features into a mask through elementwise sum
                    for mm, value in temp_dict.items():
                        curr_mask += value * avai_dict[mm].unsqueeze(1)
                        mask_num += avai_dict[mm]
                    curr_mask = torch.div(curr_mask, (mask_num + 0.00000000001).unsqueeze(1))
                else:
                    raise ValueError(f"Unknown CE mechanism: {self.use_ce}")
                curr_mask = self.h_reason(curr_mask)
                experts[curr_modality] = l(experts[curr_modality], curr_mask)

        elif self.concat_mix_experts:
            concatenated = torch.cat(tuple(experts.values()), dim=1)
            vid_embd_ = self.video_GU[0](concatenated)
            text_embd_ = text_embd[self.modalities[0]]
            text_embd_ = text_embd_.view(-1, text_embd_.shape[-1])

        elif self.concat_experts:
            vid_embd_ = torch.cat(tuple(experts.values()), dim=1)
            text_embd_ = text_embd[self.modalities[0]]
            text_embd_ = text_embd_.view(-1, text_embd_.shape[-1])
        else:
            for modality, layer in zip(self.modalities, self.video_GU):
                experts[modality] = layer(experts[modality])

        if self.training:
            merge_caption_similiarities = "avg"
        else:
            merge_caption_similiarities = self.test_caption_mode

        if self.task == "classification":
            # for modality, layer in zip(self.modalities, self.video_dim_reduce_later):
            #     attempt to perform affordable classifier, might be removed later
            #     experts[modality] = layer(experts[modality])
            concatenated = torch.cat(tuple(experts.values()), dim=1)
            preds = self.classifier(concatenated)
            return {"modalities": self.modalities, "class_preds": preds}
        elif self.concat_experts or self.concat_mix_experts:
            # zero pad to accommodate mismatch in sizes (after first setting the number
            # of VLAD clusters for the text to get the two vectors as close as possible
            # in size)
            if text_embd_.shape[1] > vid_embd_.shape[1]:
                sz = (vid_embd_.shape[0], text_embd_.shape[1])
                dtype, device = text_embd_.dtype, text_embd_.device
                vid_embd_padded = torch.zeros(size=sz, dtype=dtype, device=device)
                # try:
                #     vid_embd_padded[:, :vid_embd_.shape[1]] = vid_embd_
                # except:
                #     import ipdb; ipdb.set_trace()
                vid_embd_ = vid_embd_padded
            else:
                sz = (text_embd_.shape[0], vid_embd_.shape[1])
                dtype, device = text_embd_.dtype, text_embd_.device
                text_embd_padded = torch.zeros(size=sz, dtype=dtype, device=device)
                text_embd_padded[:, :text_embd_.shape[1]] = text_embd_
                text_embd_ = text_embd_padded
            cross_view_conf_matrix = torch.matmul(text_embd_, vid_embd_.t())
        elif self.task == "compute_video_embeddings":
            return {"modalities": self.modalities, "embeddings": experts}
        else:
            cross_view_conf_matrix = sharded_cross_view_inner_product(
                ind=ind,
                vid_embds=experts,
                text_embds=text_embd,
                keep_missing_modalities=self.keep_missing_modalities,
                l2renorm=self.l2renorm,
                text_weights=moe_weights,
                subspaces=self.modalities,
                raw_captions=raw_captions,
                merge_caption_similiarities=merge_caption_similiarities,
            )
        return {
            "modalities": self.modalities,
            "cross_view_conf_matrix": cross_view_conf_matrix,
            "text_embds": text_embd,
            "vid_embds": experts,
        }
Compute joint embeddings and, if requested, a confusion matrix between
video and text representations in the minibatch.

Notation: B = batch size, M = number of modalities
towhee/models/collaborative_experts/collaborative_experts.py