import torch


def l2norm_batch(Xr, Yr):
    """Compute \|X - Y\|_2^2 of tensors X, Y (batched version)
    Return:
        dnorm2: n x n x b
    """

    Xr = Xr.permute(1, 0, 2) # b x n x dim
    Yr = Yr.permute(1, 0, 2) # b x n x dim
    
    XrXr = torch.sum(Xr * Xr, dim=-1).unsqueeze(-1) # b x n x 1
    YrYr = torch.sum(Yr * Yr, dim=-1).unsqueeze(-2) # b x 1 x m
    XrYr = torch.matmul(Yr, torch.transpose(Xr, -1, -2))
    
    dnorm2 = XrXr + YrYr - 2 * XrYr # b x n x m
    dnorm2 = dnorm2.permute(1, 2, 0) # n x m x b

    return dnorm2