def sharded_cross_view_inner_product(vid_embds, text_embds, text_weights,
                                     subspaces, l2renorm, ind,
                                     keep_missing_modalities,
                                     merge_caption_similiarities="avg", tol=1E-5,
                                     raw_captions=None):
    """
    Compute a similarity matrix from sharded vectors.

    Args:
        embds1 (`dict`):
            The set of sub-embeddings that, when concatenated, form the whole.
            The ith shard has shape `B x K x F_i` (i.e. they can differ in
            the last dimension).
        embds2 (`dict`):
            Same format.
        weights2 (`torch.Tensor`):
            Weights for the shards in `embds2`.
        l2norm (`bool`):
            Whether to l2 renormalize the full embeddings.

    Returns:
        (`torch.Tensor`):
            Similarity matrix of size `BK x BK`.

    NOTE: If multiple captions are provided, we can aggregate their similarities to
    provide a single video-text similarity score.
    """
    _ = raw_captions
    bb = vid_embds[subspaces[0]].size(0)
    tt, num_caps, _ = text_embds[subspaces[0]].size()
    device = vid_embds[subspaces[0]].device

    # unroll separate captions onto first dimension and treat them separately
    sims = torch.zeros(tt * num_caps, bb, device=device)
    text_weights = text_weights.view(tt * num_caps, -1)

    if keep_missing_modalities:
        # assign every expert/text inner product the same weight, even if the expert
        # is missing
        text_weight_tensor = torch.ones(tt * num_caps, bb, len(subspaces),
                                        dtype=text_weights.dtype,
                                        device=text_weights.device)
    else:
        # mark expert availabilities along the second axis
        available = torch.ones(1, bb, len(subspaces), dtype=text_weights.dtype)
        for ii, modality in enumerate(subspaces):
            available[:, :, ii] = ind[modality]
        available = available.to(text_weights.device)
        msg = "expected `available` modality mask to only contain 0s or 1s"
        assert set(torch.unique(available).cpu().numpy()).issubset(set([0, 1])), msg
        # set the text weights along the first axis and combine with availabilities to
        # produce a <T x B x num_experts> tensor
        text_weight_tensor = text_weights.view(tt * num_caps, 1, len(subspaces)) * available
        # normalise to account for missing experts
        normalising_weights = text_weight_tensor.sum(2).view(tt * num_caps, bb, 1)
        text_weight_tensor = torch.div(text_weight_tensor, normalising_weights)

    if l2renorm:
        raise NotImplementedError("Do not use renorm until availability fix is complete")
    else:
        l2_mass_text, l2_mass_vid = 1, 1

    for idx, modality in enumerate(subspaces):
        vid_embd_ = vid_embds[modality].reshape(bb, -1) / l2_mass_vid
        text_embd_ = text_embds[modality].view(tt * num_caps, -1)
        msg = "expected weights to be applied to text embeddings"
        assert text_embd_.shape[0] == text_weights.shape[0], msg
        text_embd_ = text_embd_ / l2_mass_text
        weighting = text_weight_tensor[:, :, idx]
        sims += weighting * torch.matmul(text_embd_, vid_embd_.t())  # (T x num_caps) x (B)

    if l2renorm:
        # if not (sims.max() < 1 + tol):
        #     import ipdb; ipdb.set_trace()
        assert sims.max() < 1 + tol, "expected cosine similarities to be < 1"
        assert sims.min() > -1 - tol, "expected cosine similarities to be > -1"

    if torch.isnan(sims).sum().item():
        raise ValueError("Found nans in similarity matrix!")

    if num_caps > 1:
        # aggregate similarities from different captions
        if merge_caption_similiarities == "avg":
            sims = sims.view(bb, num_caps, bb)

            sims = torch.mean(sims, dim=1)
            sims = sims.view(bb, bb)
        elif merge_caption_similiarities == "indep":
            pass
        else:
            msg = "unrecognised merge mode: {}"
            raise ValueError(msg.format(merge_caption_similiarities))
    return sims
Compute a similarity matrix from sharded vectors.

Args:
embds1 (`dict`):
The set of sub-embeddings that, when concatenated, form the whole.
The ith shard has shape `B x K x F_i` (i.e. they can differ in
the last dimension).
embds2 (`dict`):
Same format.
weights2 (`torch.Tensor`):
Weights for the shards in `embds2`.
l2norm (`bool`):
Whether to l2 renormalize the full embeddings.

Returns:
(`torch.Tensor`):
Similarity matrix of size `BK x BK`.

NOTE: If multiple captions are provided, we can aggregate their similarities to
provide a single video-text similarity score.
towhee/models/collaborative_experts/collaborative_experts.py