import torch


def pdist(X, Y):
    """ Computes all the pairwise distances

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

    Returns
    -------
    torch.tensor
        shape [n, m] of all pairwise distances
    """
    n, m = X.shape[0], Y.shape[0]
    X_norm2 = (X ** 2).sum(1)
    Y_norm2 = (Y ** 2).sum(1)
    X_dot_Y = X @ Y.T

    return (
        X_norm2.unsqueeze(1) @ torch.ones((1, m), device=X.device)
        - 2 * X_dot_Y
        + torch.ones((n, 1), device=Y.device) @ Y_norm2.unsqueeze(0)
    )