def sinkhorn_knopp(a, b, C, reg=0.1, maxIter=1000, stopThr=1e-09, verbose=True, log=True, warm_start=None, eval_freq=10, print_freq=200, **kwargs):
    '\n    Solve the entropic regularization optimal transport\n    The input should be PyTorch tensors\n    The function solves the following optimization problem:\n\n    .. math::\n        \\gamma = arg\\min_\\gamma <\\gamma,C>_F + reg\\cdot\\Omega(\\gamma)\n        s.t. \\gamma 1 = a\n             \\gamma^T 1= b\n             \\gamma\\geq 0\n    where :\n    - C is the (ns,nt) metric cost matrix\n    - :math:`\\Omega` is the entropic regularization term :math:`\\Omega(\\gamma)=\\sum_{i,j} \\gamma_{i,j}\\log(\\gamma_{i,j})`\n    - a and b are target and source measures (sum to 1)\n    The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [1].\n\n    Parameters\n    ----------\n    a : torch.tensor (na,)\n        samples measure in the target domain\n    b : torch.tensor (nb,)\n        samples in the source domain\n    C : torch.tensor (na,nb)\n        loss matrix\n    reg : float\n        Regularization term > 0\n    maxIter : int, optional\n        Max number of iterations\n    stopThr : float, optional\n        Stop threshol on error ( > 0 )\n    verbose : bool, optional\n        Print information along iterations\n    log : bool, optional\n        record log if True\n\n    Returns\n    -------\n    gamma : (na x nb) torch.tensor\n        Optimal transportation matrix for the given parameters\n    log : dict\n        log dictionary return only if log==True in parameters\n\n    References\n    ----------\n    [1] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013\n    See Also\n    --------\n\n    '
    device = a.device
    (na, nb) = C.shape
    assert ((na >= 1) and (nb >= 1)), 'C needs to be 2d'
    assert ((na == a.shape[0]) and (nb == b.shape[0])), "Shape of a or b does't match that of C"
    assert (reg > 0), 'reg should be greater than 0'
    assert ((a.min() >= 0.0) and (b.min() >= 0.0)), 'Elements in a or b less than 0'
    if log:
        log = {'err': []}
    if (warm_start is not None):
        u = warm_start['u']
        v = warm_start['v']
    else:
        u = (torch.ones(na, dtype=a.dtype).to(device) / na)
        v = (torch.ones(nb, dtype=b.dtype).to(device) / nb)
    K = torch.empty(C.shape, dtype=C.dtype).to(device)
    torch.div(C, (- reg), out=K)
    torch.exp(K, out=K)
    b_hat = torch.empty(b.shape, dtype=C.dtype).to(device)
    it = 1
    err = 1
    KTu = torch.empty(v.shape, dtype=v.dtype).to(device)
    Kv = torch.empty(u.shape, dtype=u.dtype).to(device)
    while ((err > stopThr) and (it <= maxIter)):
        (upre, vpre) = (u, v)
        torch.matmul(u, K, out=KTu)
        v = torch.div(b, (KTu + M_EPS))
        torch.matmul(K, v, out=Kv)
        u = torch.div(a, (Kv + M_EPS))
        if (torch.any(torch.isnan(u)) or torch.any(torch.isnan(v)) or torch.any(torch.isinf(u)) or torch.any(torch.isinf(v))):
            print('Warning: numerical errors at iteration', it)
            (u, v) = (upre, vpre)
            break
        if (log and ((it % eval_freq) == 0)):
            b_hat = (torch.matmul(u, K) * v)
            err = (b - b_hat).pow(2).sum().item()
            log['err'].append(err)
        if (verbose and ((it % print_freq) == 0)):
            print('iteration {:5d}, constraint error {:5e}'.format(it, err))
        it += 1
    if log:
        log['u'] = u
        log['v'] = v
        log['alpha'] = (reg * torch.log((u + M_EPS)))
        log['beta'] = (reg * torch.log((v + M_EPS)))
    P = ((u.reshape((- 1), 1) * K) * v.reshape(1, (- 1)))
    if log:
        return (P, log)
    else:
        return P
Solve the entropic regularization optimal transport
The input should be PyTorch tensors
The function solves the following optimization problem:

.. math::
    \gamma = arg\min_\gamma <\gamma,C>_F + reg\cdot\Omega(\gamma)
    s.t. \gamma 1 = a
         \gamma^T 1= b
         \gamma\geq 0
where :
- C is the (ns,nt) metric cost matrix
- :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
- a and b are target and source measures (sum to 1)
The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [1].

Parameters
----------
a : torch.tensor (na,)
    samples measure in the target domain
b : torch.tensor (nb,)
    samples in the source domain
C : torch.tensor (na,nb)
    loss matrix
reg : float
    Regularization term > 0
maxIter : int, optional
    Max number of iterations
stopThr : float, optional
    Stop threshol on error ( > 0 )
verbose : bool, optional
    Print information along iterations
log : bool, optional
    record log if True

Returns
-------
gamma : (na x nb) torch.tensor
    Optimal transportation matrix for the given parameters
log : dict
    log dictionary return only if log==True in parameters

References
----------
[1] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013
See Also
--------
def sinkhorn_knopp(a, b, C, reg=0.1, maxIter=1000, stopThr=1e-09, verbose=True, log=True, warm_start=None, eval_freq=10, print_freq=200, **kwargs):
    '\n    Solve the entropic regularization optimal transport\n    The input should be PyTorch tensors\n    The function solves the following optimization problem:\n\n    .. math::\n        \\gamma = arg\\min_\\gamma <\\gamma,C>_F + reg\\cdot\\Omega(\\gamma)\n        s.t. \\gamma 1 = a\n             \\gamma^T 1= b\n             \\gamma\\geq 0\n    where :\n    - C is the (ns,nt) metric cost matrix\n    - :math:`\\Omega` is the entropic regularization term :math:`\\Omega(\\gamma)=\\sum_{i,j} \\gamma_{i,j}\\log(\\gamma_{i,j})`\n    - a and b are target and source measures (sum to 1)\n    The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [1].\n\n    Parameters\n    ----------\n    a : torch.tensor (na,)\n        samples measure in the target domain\n    b : torch.tensor (nb,)\n        samples in the source domain\n    C : torch.tensor (na,nb)\n        loss matrix\n    reg : float\n        Regularization term > 0\n    maxIter : int, optional\n        Max number of iterations\n    stopThr : float, optional\n        Stop threshol on error ( > 0 )\n    verbose : bool, optional\n        Print information along iterations\n    log : bool, optional\n        record log if True\n\n    Returns\n    -------\n    gamma : (na x nb) torch.tensor\n        Optimal transportation matrix for the given parameters\n    log : dict\n        log dictionary return only if log==True in parameters\n\n    References\n    ----------\n    [1] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013\n    See Also\n    --------\n\n    '
    device = a.device
    (na, nb) = C.shape
    assert ((na >= 1) and (nb >= 1)), 'C needs to be 2d'
    assert ((na == a.shape[0]) and (nb == b.shape[0])), "Shape of a or b does't match that of C"
    assert (reg > 0), 'reg should be greater than 0'
    assert ((a.min() >= 0.0) and (b.min() >= 0.0)), 'Elements in a or b less than 0'
    if log:
        log = {'err': []}
    if (warm_start is not None):
        u = warm_start['u']
        v = warm_start['v']
    else:
        u = (torch.ones(na, dtype=a.dtype).to(device) / na)
        v = (torch.ones(nb, dtype=b.dtype).to(device) / nb)
    K = torch.empty(C.shape, dtype=C.dtype).to(device)
    torch.div(C, (- reg), out=K)
    torch.exp(K, out=K)
    b_hat = torch.empty(b.shape, dtype=C.dtype).to(device)
    it = 1
    err = 1
    KTu = torch.empty(v.shape, dtype=v.dtype).to(device)
    Kv = torch.empty(u.shape, dtype=u.dtype).to(device)
    while ((err > stopThr) and (it <= maxIter)):
        (upre, vpre) = (u, v)
        torch.matmul(u, K, out=KTu)
        v = torch.div(b, (KTu + M_EPS))
        torch.matmul(K, v, out=Kv)
        u = torch.div(a, (Kv + M_EPS))
        if (torch.any(torch.isnan(u)) or torch.any(torch.isnan(v)) or torch.any(torch.isinf(u)) or torch.any(torch.isinf(v))):
            print('Warning: numerical errors at iteration', it)
            (u, v) = (upre, vpre)
            break
        if (log and ((it % eval_freq) == 0)):
            b_hat = (torch.matmul(u, K) * v)
            err = (b - b_hat).pow(2).sum().item()
            log['err'].append(err)
        if (verbose and ((it % print_freq) == 0)):
            print('iteration {:5d}, constraint error {:5e}'.format(it, err))
        it += 1
    if log:
        log['u'] = u
        log['v'] = v
        log['alpha'] = (reg * torch.log((u + M_EPS)))
        log['beta'] = (reg * torch.log((v + M_EPS)))
    P = ((u.reshape((- 1), 1) * K) * v.reshape(1, (- 1)))
    if log:
        return (P, log)
    else:
        return P