import torch


def calc_dist_mat(a_coord: torch.Tensor,
                  b_coord: torch.Tensor) -> torch.Tensor:
    """
    Calculate a distance matrix between tensors of coords
    """
    assert a_coord.shape == b_coord.shape
    mat_shape = (len(a_coord), len(a_coord), 3)

    a_coord = a_coord.unsqueeze(0).expand(mat_shape)
    b_coord = b_coord.unsqueeze(1).expand(mat_shape)

    dist_mat = (a_coord - b_coord).norm(dim=-1)

    return dist_mat