import torch


def cartesian_product(a, b):
    """
    Compute the batched cartesian product between two matrices.
    Input:
        a: Tensor(n, d1)
        b: Tensor(n, d2)
    Output:
        output: Tensor(n, d1 * d2, 2)
    """
    n1, d1 = a.shape
    n2, d2 = b.shape
    assert n1 == n2
    return torch.cat([
        a.unsqueeze(-1).repeat(1, 1, d2).unsqueeze(-1),
        b.repeat(1, d1).view(n2, d1, d2).unsqueeze(-1)
    ], 3).view(n1, d1 * d2, 2)