import torch


def concat_per_row(A, B):
    """Concat every row in A with every row in B"""
    m1, n1 = A.shape
    m2, n2 = B.shape

    res = torch.zeros(m1, m2, n1 + n2)
    res[:, :, :n1] = A[:, None, :]
    res[:, :, n1:] = B
    return res.view(m1 * m2, -1)