def sharded_single_view_inner_product(embds, subspaces, text_weights=None,
                                      l2renorm=True):
    """
    Compute a similarity matrix from sharded vectors.

    Args:
        embds (`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), or shape `B x F_i`
        l2norm (`bool`):
            Whether to l2 normalize the full embedding.

    Returns:
        (`torch.Tensor`):
            Similarity matrix of size `BK x BK`.
    """
    _ = subspaces
    subspaces = list(embds.keys())
    device = embds[subspaces[0]].device
    shape = embds[subspaces[0]].shape
    if len(shape) == 3:
        bb, kk, _ = shape
        num_embds = bb * kk
        assert text_weights is not None, "Expected 3-dim tensors for text (+ weights)"
        assert text_weights.shape[0] == bb
        assert text_weights.shape[1] == kk
    elif len(shape) == 2:
        bb, _ = shape
        num_embds = bb
        assert text_weights is None, "Expected 2-dim tensors for non-text (no weights)"
    else:
        raise ValueError("input tensor with {} dims unrecognised".format(len(shape)))
    sims = torch.zeros(num_embds, num_embds, device=device)
    if l2renorm:
        l2_mass = 0
        for idx, modality in enumerate(subspaces):
            embd_ = embds[modality]
            if text_weights is not None:
                # text_weights (i.e. moe_weights) are shared among subspace for video
                embd_ = text_weights[:, :, idx:idx + 1] * embd_
            embd_ = embds[modality].reshape(num_embds, -1)
            l2_mass += embd_.pow(2).sum(1)
        l2_mass = torch.sqrt(l2_mass.clamp(min=1E-6)).unsqueeze(1)
    else:
        l2_mass = 1

    for idx, modality in enumerate(subspaces):
        embd_ = embds[modality]
        if text_weights is not None:
            embd_ = text_weights[:, :, idx:idx + 1] * embd_
        embd_ = embd_.reshape(num_embds, -1) / l2_mass
        sims += torch.matmul(embd_, embd_.t())
    if torch.isnan(sims).sum().item():
        raise ValueError("Found nans in similarity matrix!")
    return sims
Compute a similarity matrix from sharded vectors.

Args:
embds (`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), or shape `B x F_i`
l2norm (`bool`):
Whether to l2 normalize the full embedding.

Returns:
(`torch.Tensor`):
Similarity matrix of size `BK x BK`.
towhee/models/collaborative_experts/collaborative_experts.py