def sinkhorn_epsilon_scaling(a, b, C, reg=0.1, maxIter=100, maxInnerIter=100, tau=1000.0, scaling_base=0.75, scaling_coef=None, stopThr=1e-09, verbose=False, log=False, warm_start=None, eval_freq=10, print_freq=200, **kwargs):
    '\n    Solve the entropic regularization OT problem with log stabilization\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\n    The algorithm used for solving the problem is the Sinkhorn-Knopp matrix\n    scaling algorithm as proposed in [1] but with the log stabilization\n    proposed in [3] and the log scaling proposed in [2] algorithm 3.2\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    tau : float\n        thershold for max value in u or v for log scaling\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    [2] Bernhard Schmitzer. Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems. SIAM Journal on Scientific Computing, 2019\n    [3] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.\n\n    See Also\n    --------\n\n    '
    (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'

    def get_reg(it, reg, pre_reg):
        if (it == 1):
            return scaling_coef
        elif (((pre_reg - reg) * scaling_base) < M_EPS):
            return reg
        else:
            return (((pre_reg - reg) * scaling_base) + reg)
    if (scaling_coef is None):
        scaling_coef = (C.max() + reg)
    it = 1
    err = 1
    running_reg = scaling_coef
    if log:
        log = {'err': []}
    warm_start = None
    while ((err > stopThr) and (it <= maxIter)):
        running_reg = get_reg(it, reg, running_reg)
        (P, _log) = sinkhorn_stabilized(a, b, C, running_reg, maxIter=maxInnerIter, tau=tau, stopThr=stopThr, verbose=False, log=True, warm_start=warm_start, eval_freq=eval_freq, print_freq=print_freq, **kwargs)
        warm_start = {}
        warm_start['alpha'] = _log['alpha']
        warm_start['beta'] = _log['beta']
        primal_val = (((C * P).sum() + (reg * (P * torch.log(P)).sum())) - (reg * P.sum()))
        dual_val = (((_log['alpha'] * a).sum() + (_log['beta'] * b).sum()) - (reg * P.sum()))
        err = (primal_val - dual_val)
        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['alpha'] = _log['alpha']
        log['beta'] = _log['beta']
        return (P, log)
    else:
        return P
Solve the entropic regularization OT problem with log stabilization
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] but with the log stabilization
proposed in [3] and the log scaling proposed in [2] algorithm 3.2

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
tau : float
    thershold for max value in u or v for log scaling
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
[2] Bernhard Schmitzer. Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems. SIAM Journal on Scientific Computing, 2019
[3] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.

See Also
--------
def sinkhorn_epsilon_scaling(a, b, C, reg=0.1, maxIter=100, maxInnerIter=100, tau=1000.0, scaling_base=0.75, scaling_coef=None, stopThr=1e-09, verbose=False, log=False, warm_start=None, eval_freq=10, print_freq=200, **kwargs):
    '\n    Solve the entropic regularization OT problem with log stabilization\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\n    The algorithm used for solving the problem is the Sinkhorn-Knopp matrix\n    scaling algorithm as proposed in [1] but with the log stabilization\n    proposed in [3] and the log scaling proposed in [2] algorithm 3.2\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    tau : float\n        thershold for max value in u or v for log scaling\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    [2] Bernhard Schmitzer. Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems. SIAM Journal on Scientific Computing, 2019\n    [3] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.\n\n    See Also\n    --------\n\n    '
    (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'

    def get_reg(it, reg, pre_reg):
        if (it == 1):
            return scaling_coef
        elif (((pre_reg - reg) * scaling_base) < M_EPS):
            return reg
        else:
            return (((pre_reg - reg) * scaling_base) + reg)
    if (scaling_coef is None):
        scaling_coef = (C.max() + reg)
    it = 1
    err = 1
    running_reg = scaling_coef
    if log:
        log = {'err': []}
    warm_start = None
    while ((err > stopThr) and (it <= maxIter)):
        running_reg = get_reg(it, reg, running_reg)
        (P, _log) = sinkhorn_stabilized(a, b, C, running_reg, maxIter=maxInnerIter, tau=tau, stopThr=stopThr, verbose=False, log=True, warm_start=warm_start, eval_freq=eval_freq, print_freq=print_freq, **kwargs)
        warm_start = {}
        warm_start['alpha'] = _log['alpha']
        warm_start['beta'] = _log['beta']
        primal_val = (((C * P).sum() + (reg * (P * torch.log(P)).sum())) - (reg * P.sum()))
        dual_val = (((_log['alpha'] * a).sum() + (_log['beta'] * b).sum()) - (reg * P.sum()))
        err = (primal_val - dual_val)
        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['alpha'] = _log['alpha']
        log['beta'] = _log['beta']
        return (P, log)
    else:
        return P