import torch


def batch_pdist(X, Y, device=None):
    """ Computes all the pairwise distances.

    Parameters
    ----------
    X : torch.tensor
        shape [c, n, d]
    Y : torch.tensor
        shape [c, m, d] or [1, m, d]

    Returns
    -------
    torch.tensor
        output has shape [c, n, m]. The entry at i, j, k is the squared distance
        between the d-dimensional vectors X[i, j] and Y[i, k] if Y.shape[0] is c
        and Y[0, k] if Y.shape[0] is 1.
    """
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    c = X.shape[0]
    n, m = X.shape[1], Y.shape[1]
    X_norm2 = (X ** 2).sum(2)
    Y_norm2 = (Y ** 2).sum(2)
    X_dot_Y = torch.matmul(X, Y.transpose(1, 2))

    return (
        torch.matmul(X_norm2.unsqueeze(2), torch.ones((c, 1, m), device=device))
        - 2 * X_dot_Y
        + torch.matmul(torch.ones((c, n, 1), device=device), Y_norm2.unsqueeze(1))
    )