import torch


def batched_concat_per_row(A, B):
    """Concat every row in A with every row in B where
    the first dimension of A and B is the batch dimension"""
    b, m1, n1 = A.shape
    _, m2, n2 = B.shape

    res = torch.zeros(b, m1, m2, n1 + n2)
    res[:, :, :, :n1] = A[:, :, None, :]

    res[:, :, :, n1:] = B[:, None, :, :]
    return res.view(b, m1 * m2, -1)