How to implement pair-wise calculation of attention within a batch?

Suppose I now have the following code to calculate source-target attention for two variable, x and y:

class MultiHeadedAttention(nn.Module):
    """Multi-Head Attention layer

    :param int n_head: the number of head s
    :param int n_feat: the number of features
    :param float dropout_rate: dropout rate
    """

    def __init__(self, n_head: int, n_feat: int, dropout_rate: float):
        super(MultiHeadedAttention, self).__init__()
        assert n_feat % n_head == 0
        self.d_k = n_feat // n_head
        self.h = n_head
        self.linear_q = nn.Linear(n_feat, n_feat)
        self.linear_k = nn.Linear(n_feat, n_feat)
        self.linear_v = nn.Linear(n_feat, n_feat)
        self.linear_out = nn.Linear(n_feat, n_feat)
        self.dropout = nn.Dropout(p=dropout_rate)

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Compute 'Scaled Dot Product Attention'

        :param torch.Tensor query: (batch, x_len, size)
        :param torch.Tensor key: (batch, y_len, size)
        :param torch.Tensor value: (batch, y_len, size)
        :param torch.Tensor mask: (batch, x_len, y_len)
        :param torch.nn.Dropout dropout:
        :return torch.Tensor: attentined and transformed `value` (batch, x_len, depth)
             weighted by the query dot key attention (batch, head, x_len, y_len)
        """
        n_batch = query.size(0)
        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
        q = q.transpose(1, 2)  # (batch, head, x_len, d_k)
        k = k.transpose(1, 2)  # (batch, head, x_len, d_k)
        v = v.transpose(1, 2)  # (batch, head, y_len, d_k)
        
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(
            self.d_k
        )  # (batch, head, x_len, y_len)
        if mask is not None:
            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, x_len, y_len)
            mask = mask.to(device=scores.device)
            scores = scores.masked_fill_(mask, -np.inf)
            attn = torch.softmax(scores, dim=-1).masked_fill(
                mask, 0.0
            )  # (batch, head, x_len, y_len)
        else:
            attn = torch.softmax(scores, dim=-1)  # (batch, head, x_len, y_len)

        p_attn = self.dropout(attn)
        x = torch.matmul(p_attn, v)  # (batch, head, x_len, d_k)
        x = (
            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
        )  # (batch, x_len, depth)
        return self.linear_out(x)  # (batch, x_len, depth)

So this class calculate the attention of batch size=B pairs of (x, y)_i, gives output of dim (batch, x_len, depth). So far so good.

The question is: What if I wanted to extend this class to calculate NOT ONLY (x1, y1), (x2, y2)..., but also all combination of xy, i.e. (x1, y2), (x1, y3)... within the batch, so that I will get an output of dim (batch, batch, x_len, depth) WITHOUT LOOPING.

How would you implement this? Any recommendation, suggestion, example is appreciated.

EDITED

I just came up with an idea which does the desired job at the expense of extra memory use. Just simply copy X and Y along the batch dimension so that the represent all the pairs of x_i and y_i. Specifically:

b = torch.tensor(list(range(batch_size)))
comb = torch.cartesian_prod(b, b)
x = x[comb[:, 0], :, :]
y = y[comb[:, 1], :, :]

and then after the calculation, view or reshape the first dimension and it will return output which is of dim=(batch_size, batch_size, x_len, depth).

I have tested using toy example and quite sure it does do the job. However, unfortunately, for my case it got CUDA out of memory.
What would you do under this situation? Should I give up on parallelism and just use loop to make it works?