import torch


def fix_coupling_sinkhorn(log_coupling, log_p1, log_p2, iterations=10):
    """Adjust a coupling to approximately match marginals using Sinkhorn iteration.
    Args:
    log_coupling: Log of a coupling, possibly unnormalized, of shape [M,M]
    log_p1: Log probabilities for first dimension logits, of shape [M].
    log_p2: Log probabilities for second dimension logits, of shape [M].
    iterations: How many Sinkhorn iterations to use.
    Returns:
    Matrix of shape [M,M]. Will always exactly match the marginals of `log_p2`
    along the second axis, and will attempt to match marginals of `log_p1` along
    the first axis as well, but may not reach it depending on iteration count.
    """
    for _ in range(iterations):
        log_coupling = log_coupling + log_p1[:, :, None] - torch.logsumexp(log_coupling, dim=2, keepdim=True)
        log_coupling = log_coupling + log_p2[:, None, :] - torch.logsumexp(log_coupling, dim=1, keepdim=True)
    return log_coupling