import torch


def _pairwise_squared_distance_matrix(x:torch.Tensor) -> torch.Tensor:
    """Pairwise squared distance among a (batch) matrix's rows (2nd dim).

    This saves a bit of computation vs. using
    `_cross_squared_distance_matrix(x, x)`

    Args:
        x: torch.Tensor [shape=(b, n, d), float]
    Returns:
        squared_dists: torch.Tensor [shape=(b, n, n), float]
            `squared_dists[b,i,j] = ||x[b,i,:] - x[b,j,:]||^2`.
    """
    x_x_transpose = torch.einsum('bnd, bmd -> bnm', x, x)  # shape=(b, n, n)
    x_norm_squared = torch.diagonal(x_x_transpose, dim1=1, dim2=2).unsqueeze(2)  # shape=(b, n, 1)
    # squared_dists[b,i,j] = ||x_bi - x_bj||^2 = x_bi'x_bi- 2x_bi'x_bj + x_bj'x_bj
    squared_dists = x_norm_squared - 2*x_x_transpose + x_norm_squared.transpose(1, 2)
    return squared_dists