Compute the sgd algorithm to solve the regularized discrete measures
        optimal transport dual problem

    The function solves the following optimization problem:

    .. math::
        \gamma = arg\min_\gamma <\gamma,M>_F + reg\cdot\Omega(\gamma)

        s.t. \gamma 1 = a

             \gamma^T 1= b

             \gamma \geq 0

    Where :

    - M is the (ns,nt) metric cost matrix
    - :math:`\Omega` is the entropic regularization term with :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
    - a and b are source and target weights (sum to 1)

    Parameters
    ----------

    a : np.ndarray(ns,)
        source measure
    b : np.ndarray(nt,)
        target measure
    M : np.ndarray(ns, nt)
        cost matrix
    reg : float number
        Regularization term > 0
    batch_size : int number
        size of the batch
    numItermax : int number
        number of iteration
    lr : float number
        learning rate

    Returns
    -------

    alpha : np.ndarray(ns,)
        dual variable
    beta : np.ndarray(nt,)
        dual variable

    Examples
    --------

    >>> n_source = 7
    >>> n_target = 4
    >>> reg = 1
    >>> numItermax = 20000
    >>> lr = 0.1
    >>> batch_size = 3
    >>> log = True
    >>> a = ot.utils.unif(n_source)
    >>> b = ot.utils.unif(n_target)
    >>> rng = np.random.RandomState(0)
    >>> X_source = rng.randn(n_source, 2)
    >>> Y_target = rng.randn(n_target, 2)
    >>> M = ot.dist(X_source, Y_target)
    >>> sgd_dual_pi, log = stochastic.solve_dual_entropic(a, b, M, reg,
                                                          batch_size,
                                                          numItermax, lr, log)
    >>> print(log['alpha'], log['beta'])
    >>> print(sgd_dual_pi)

    References
    ----------

    [Seguy et al., 2018] :
                    International Conference on Learning Representation (2018),
                      arXiv preprint arxiv:1711.02283.

def sgd_entropic_regularization(a, b, M, reg, batch_size, numItermax, lr):
    '''
    Compute the sgd algorithm to solve the regularized discrete measures
        optimal transport dual problem

    The function solves the following optimization problem:

    .. math::
        \gamma = arg\min_\gamma <\gamma,M>_F + reg\cdot\Omega(\gamma)

        s.t. \gamma 1 = a

             \gamma^T 1= b

             \gamma \geq 0

    Where :

    - M is the (ns,nt) metric cost matrix
    - :math:`\Omega` is the entropic regularization term with :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
    - a and b are source and target weights (sum to 1)

    Parameters
    ----------

    a : np.ndarray(ns,)
        source measure
    b : np.ndarray(nt,)
        target measure
    M : np.ndarray(ns, nt)
        cost matrix
    reg : float number
        Regularization term > 0
    batch_size : int number
        size of the batch
    numItermax : int number
        number of iteration
    lr : float number
        learning rate

    Returns
    -------

    alpha : np.ndarray(ns,)
        dual variable
    beta : np.ndarray(nt,)
        dual variable

    Examples
    --------

    >>> n_source = 7
    >>> n_target = 4
    >>> reg = 1
    >>> numItermax = 20000
    >>> lr = 0.1
    >>> batch_size = 3
    >>> log = True
    >>> a = ot.utils.unif(n_source)
    >>> b = ot.utils.unif(n_target)
    >>> rng = np.random.RandomState(0)
    >>> X_source = rng.randn(n_source, 2)
    >>> Y_target = rng.randn(n_target, 2)
    >>> M = ot.dist(X_source, Y_target)
    >>> sgd_dual_pi, log = stochastic.solve_dual_entropic(a, b, M, reg,
                                                          batch_size,
                                                          numItermax, lr, log)
    >>> print(log['alpha'], log['beta'])
    >>> print(sgd_dual_pi)

    References
    ----------

    [Seguy et al., 2018] :
                    International Conference on Learning Representation (2018),
                      arXiv preprint arxiv:1711.02283.
    '''

    n_source = np.shape(M)[0]
    n_target = np.shape(M)[1]
    cur_alpha = np.zeros(n_source)
    cur_beta = np.zeros(n_target)
    for cur_iter in range(numItermax):
        k = np.sqrt(cur_iter + 1)
        batch_alpha = np.random.choice(n_source, batch_size, replace=False)
        batch_beta = np.random.choice(n_target, batch_size, replace=False)
        update_alpha, update_beta = batch_grad_dual(a, b, M, reg, cur_alpha,
                                                    cur_beta, batch_size,
                                                    batch_alpha, batch_beta)
        cur_alpha[batch_alpha] += (lr / k) * update_alpha[batch_alpha]
        cur_beta[batch_beta] += (lr / k) * update_beta[batch_beta]

    return cur_alpha, cur_beta