//  [python]
def UpdatePlans_LSE(X, Y, Z, a, b, reg, cost, max_iter=1000, delta=1e-9, lam=0):

    C1 = cost(Z, X)
    C2 = cost(Z, Y)

    r = np.shape(Z)[0]
    f1, f2 = np.zeros(r), np.zeros(r)
    g1, g2 = np.zeros(np.shape(a)[0]), np.zeros(np.shape(b)[0])

    w = np.ones(r) / r

    err = 1
    n_iter = 0
    while n_iter < max_iter:
        f1_prev, g1_prev = f1, g1
        f2_prev, g2_prev = f2, g2
        w_prev = w
        if err > delta:
            n_iter = n_iter + 1

            # Update g1, g2
            C1_tilde = (
                f1.reshape(-1, 1) * np.ones((1, np.shape(a)[0]))
                + np.ones((r, 1)) * g1.reshape(1, -1)
                - C1
            )
            C1_tilde = C1_tilde / reg
            g1 = reg * np.log(a) + g1 - reg * np.max(C1_tilde, axis=0) - reg * np.log(np.sum(np.exp(C1_tilde - np.max(C1_tilde, axis=0)), axis=0))

            C2_tilde = (
                f2.reshape(-1, 1) * np.ones((1, np.shape(b)[0]))
                + np.ones((r, 1)) * g2.reshape(1, -1)
                - C2
            )
            C2_tilde = C2_tilde / reg
            g2 = reg * np.log(b) + g2 - reg * np.max(C2_tilde, axis=0) - reg * np.log(np.sum(np.exp(C2_tilde - np.max(C2_tilde, axis=0)), axis=0))

            # Update w
            C1_tilde = (
                f1.reshape(-1, 1) * np.ones((1, np.shape(a)[0]))
                + np.ones((r, 1)) * g1.reshape(1, -1)
                - C1
            )
            C1_tilde = C1_tilde / reg
            P1 = np.exp(C1_tilde)

            C2_tilde = (
                f2.reshape(-1, 1) * np.ones((1, np.shape(b)[0]))
                + np.ones((r, 1)) * g2.reshape(1, -1)
                - C2
            )
            C2_tilde = C2_tilde / reg
            P2 = np.exp(C2_tilde)

            w = (np.sum(P1, axis=1) * np.sum(P2, axis=1)) ** (1 / 2)
            log_w = (1 / 2) * (
                np.max(C1_tilde, axis=1) + np.log(np.sum(np.exp(C1_tilde - np.max(C1_tilde, axis=1)), axis=1))
                + np.max(C2_tilde, axis=1) + np.log(np.sum(np.exp(C2_tilde - np.max(C2_tilde, axis=1)), axis=1))
            )

            # Update f1, f2
            C1_tilde = (
                f1.reshape(-1, 1) * np.ones((1, np.shape(a)[0]))
                + np.ones((r, 1)) * g1.reshape(1, -1)
                - C1
            )
            C1_tilde = C1_tilde / reg
            f1 = reg * log_w + f1 - reg * np.max(C1_tilde, axis=1) - reg * np.log(np.sum(np.exp(C1_tilde - np.max(C1_tilde, axis=1)), axis=1))

            C2_tilde = (
                f2.reshape(-1, 1) * np.ones((1, np.shape(b)[0]))
                + np.ones((r, 1)) * g2.reshape(1, -1)
                - C2
            )
            C2_tilde = C2_tilde / reg
            f2 = reg * log_w + f2 - reg * np.max(C2_tilde, axis=1) - reg * np.log(np.sum(np.exp(C2_tilde - np.max(C2_tilde, axis=1)), axis=1))

            # Update the coupling P1, P2
            C1_tilde = (
                f1.reshape(-1, 1) * np.ones((1, np.shape(a)[0]))
                + np.ones((r, 1)) * g1.reshape(1, -1)
                - C1
            )
            C1_tilde = C1_tilde / reg
            P1 = np.exp(C1_tilde)

            C2_tilde = (
                f2.reshape(-1, 1) * np.ones((1, np.shape(b)[0]))
                + np.ones((r, 1)) * g2.reshape(1, -1)
                - C2
            )
            C2_tilde = C2_tilde / reg
            P2 = np.exp(C2_tilde)

            # Update the error
            err_1 = np.sum(np.abs(np.sum(P1, axis=0) - a))
            err_2 = np.sum(np.abs(np.sum(P2, axis=0) - b))
            err = err_1 + err_2

            if (
                np.any(np.isnan(f1))
                or np.any(np.isnan(g1))
                or np.any(np.isnan(f2))
                or np.any(np.isnan(g2))
                or np.any(np.isinf(f1))
                or np.any(np.isinf(g1))
                or np.any(np.isinf(f2))
                or np.any(np.isinf(g2))
            ):
                # we have reached the machine precision
                # come back to previous solution and quit loop
                print("Warning: numerical errors at iteration", n_iter)
                f1, g1 = f1_prev, g1_prev
                f2, g2 = f2_prev, g2_prev
                w = w_prev
                break
        else:
            return P1, P2, w

    # Update the coupling P1, P2
    C1_tilde = (
        f1.reshape(-1, 1) * np.ones((1, np.shape(a)[0]))
        + np.ones((r, 1)) * g1.reshape(1, -1)
        - C1
    )
    C1_tilde = C1_tilde / reg
    P1 = np.exp(C1_tilde)

    C2_tilde = (
        f2.reshape(-1, 1) * np.ones((1, np.shape(b)[0]))
        + np.ones((r, 1)) * g2.reshape(1, -1)
        - C2
    )
    C2_tilde = C2_tilde / reg
    P2 = np.exp(C2_tilde)

    return P1, P2, w