mdot-tnt 0.1.0__py3-none-any.whl

mdot_tnt/__init__.py

@@ -0,0 +1,57 @@
+
+import torch as th
+
+from mdot_tnt.mdot import mdot
+from mdot_tnt.rounding import round_altschuler, rounded_cost_altschuler
+
+
+def solve_OT(r, c, C, gamma_f=4096., drop_tiny=False, return_plan=False, round=True, log=False):
+    """
+    Solve the entropic-regularized optimal transport problem. Inputs r, c, C are required to be torch tensors.
+    :param r: n-dimensional row marginal.
+    :param c: m-dimensional column marginal.
+    :param C: n x m cost matrix. Recommended use is to scale the entries to be in [0, 1].
+    :param gamma_f: The temperature (inverse of the regularization weight). For many problems, stable up to 2^18.
+    Higher values return more accurate solutions but take longer to converge. Use double precision if gamma_f large.
+    :param drop_tiny: If either marginal is known to be sparse, set this to True to drop tiny entries for a speedup.
+    If return_plan is True, the returned plan will be in the original dimensions.
+    :param return_plan: If True, return the optimal transport plan rather than the cost.
+    :param round: If True, use the rounding algorithm of Altschuler et al. (2017) to (a) return a feasible plan
+    if return_plan is True and (b) the cost of the rounded plan if return_plan is False.
+    :param log: If True, additionally return a dictionary containing logs of the optimization process.
+    :return: If return_plan is True, return the optimal transport plan as a torch tensor. Otherwise, return the cost.
+    """
+    assert all(isinstance(x, th.Tensor) for x in [r, c, C]), "r, c, and C must be torch tensors"
+    dtype = r.dtype
+    if gamma_f > 2 ** 10:
+        r, c, C = r.double(), c.double(), C.double()
+    if drop_tiny:
+        drop_lessthan = math.log(min(r.size(-1), c.size(-1))) / (gamma_f ** 2)
+        (r_, r_keep), (c_, c_keep), C_ = preprocess_marginals(r, c, C, drop_lessthan)
+
+        u_, v_, gamma_f_, k_total, opt_logs = mdot(r_, c_, C_, gamma_f)
+
+        u = -th.ones_like(r) * float('inf')
+        u[:, r_keep] = u_
+        v = -th.ones_like(c) * float('inf')
+        v[:, c_keep] = v_
+    else:
+        u, v, gamma_f_, k_total, opt_logs = mdot(r, c, C, gamma_f)
+
+    u, v = u.to(dtype), v.to(dtype)
+
+    if return_plan:
+        P = (u.unsqueeze(-1) + v.unsqueeze(-2) - gamma_f_ * C).exp()
+        if round:
+            P = round_altschuler(P, r, c)
+        if log:
+            return P, opt_logs
+        return P
+    else:
+        if round:
+            cost = rounded_cost_altschuler(u, v, r, c, C, gamma_f_)
+        else:
+            cost = ((u.unsqueeze(-1) + v.unsqueeze(-2) - gamma_f_ * C).exp() * C).sum()
+        if log:
+            return cost, opt_logs
+        return cost

mdot_tnt/mdot.py

@@ -0,0 +1,139 @@
+"""
+Code for solving the entropic-regularized optimal transport problem via the MDOT-TruncatedNewton (MDOT-TNT)
+method introduced in the paper "A Truncated Newton Method for Optimal Transport"
+by Mete Kemertas, Amir-massoud Farahmand, Allan D. Jepson (ICLR, 2025).
+URL: https://openreview.net/forum?id=gWrWUaCbMa
+"""
+
+
+import math
+import warnings
+
+from mdot_tnt.rounding import *
+from mdot_tnt.truncated_newton import TruncatedNewtonProjector
+
+
+def preprocess_marginals(r, c, C, eps):
+    """
+    This function drops the smallest entries whose cumulative sum equals 
+    :param r:
+    :param c:
+    :param C:
+    :param eps:
+    :return:
+    """
+
+    def preprocess_marginal(m, eps):
+        m_sorted, m_idx = th.sort(m, dim=-1, descending=False)
+        m_cumsum = th.cumsum(m_sorted, dim=-1)
+        m_keep = m_idx[m_cumsum > eps]
+        m_new = m[:, m_keep]
+        mass_removed = 1 - m_new.sum(-1)
+        m_new = m_new + mass_removed / m_new.size(-1)
+
+        return m_new, m_keep
+
+    r_new, r_keep = preprocess_marginal(r, eps)
+    c_new, c_keep = preprocess_marginal(c, eps)
+
+    C = C[r_keep][:, c_keep]
+
+    return (r_new, r_keep), (c_new, c_keep), C
+
+
+def smooth_marginals(r, c, eps, w_r=0.5, w_c=0.5):
+    assert w_r + w_c == 1, "w_r and w_c must sum to 1"
+    eps = eps.clamp(max=1.).unsqueeze(-1)
+    r_hat = (1 - w_r * eps) * r + w_r * eps * th.ones_like(r) / r.size(-1)
+    c_hat = (1 - w_c * eps) * c + w_c * eps * th.ones_like(c) / c.size(-1)
+
+    return r_hat, c_hat
+
+
+def adjust_schedule(q, deltas=None):
+    if deltas is None:
+        return q
+
+    deltas = deltas + [1.]  # If deltas is empty, we assume that the first iteration was successful
+    delta_min = min(deltas)
+
+    if delta_min < 0.5:
+        q = q ** 0.5
+    elif delta_min > 0.9:
+        q = q ** 2
+
+    return q
+
+
+def mdot(r, c, C, gamma_f, gamma_i=16, p=1.5, q=2.0**(1/3)):
+    """
+    Solve the entropic-regularized optimal transport problem using the MDOT method introduced in the paper:
+    "Efficient and Accurate Optimal Transport with Mirror Descent and Conjugate Gradients" by Mete Kemertas,
+    Allan D. Jepson and Amir-massoud Farahmand. URL: https://arxiv.org/abs/2307.08507
+    Here, we use the Truncated Newton method for projection.
+    :param r: The first marginal.
+    :param c: The second marginal.
+    :param C: The cost matrix. Recommended use is to scale the entries to be in [0, 1].
+    :param gamma_f: The final temperature (inverse of the regularization weight).
+    :param gamma_i: The initial temperature.
+    :param p: The exponent for the epsilon function, used to determine the stopping criterion for the dual gradient.
+    :param q: The temperature annealing (or mirror descent step size) schedule adjustment factor.
+    :return:
+    """
+    projector = TruncatedNewtonProjector(device=C.device, dtype=C.dtype)
+
+    H_r = -(r * (r + 1e-30).log()).sum(-1)
+    H_c = -(c * (c + 1e-30).log()).sum(-1)
+    H_min = th.min(H_r, H_c)
+    eps_fn = lambda g_: H_min / (g_ ** p)
+
+    logs = {
+        "proj_logs": [],
+        "eps": [],
+    }
+
+    t = 1
+    done = False
+    gamma = min(gamma_i, gamma_f)
+    gammas = [0., gamma]
+
+    while not done:
+        done = abs(gamma - gamma_f) < 1e-5  # Check if gamma == gamma_f (modulo rounding errors)
+
+        eps_d = eps_fn(gamma)
+
+        r_hat, c_hat = smooth_marginals(r, c, eps_d / 2, w_r=0.9, w_c=0.1)
+
+        if t == 1:
+            u_init, v_init = r_hat.log(), c_hat.log()
+            u_cur, v_cur = u_init.clone(), v_init.clone()
+
+        u_prev, v_prev = u_cur.clone(), v_cur.clone()
+        u_cur, v_cur, proj_log, success = projector.project(
+            gamma * C, r_hat.log(), c_hat.log(), eps_d / 2, u_init, v_init)
+
+        logs["proj_logs"].append(proj_log)
+        if not success:
+            warnings.warn("Projection failed. Returning result at the last temperature: {:.4e}".format(1 / gammas[-2]))
+            u_cur = u_prev.clone()
+            v_cur = v_prev.clone()
+            gammas = gammas[:-1]
+            break
+
+        q = adjust_schedule(q, proj_log["deltas"])
+        gamma = min(gamma * q, gamma_f)
+
+        if not done:
+            # Generate warm-started initializations for the next iteration.
+            u_init = u_cur + (u_cur - u_prev) * (gamma - gammas[-1]) / (gammas[-1] - gammas[-2])
+            v_init = v_cur + (v_cur - v_prev) * (gamma - gammas[-1]) / (gammas[-1] - gammas[-2])
+
+        gammas.append(gamma)
+        t += 1
+
+    k_total = sum([log["n_iter"] for log in logs["proj_logs"]])
+    k_total += t - 1
+    logs["success"] = success
+    logs["gammas"] = gammas
+
+    return u_cur, v_cur, gammas[-1], k_total, logs

mdot_tnt/rounding.py

@@ -0,0 +1,62 @@
+"""
+This file contains the implementation of the rounding algorithm proposed by Altschuler et al. (2017) in the paper
+"Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration". The algorithm is used to
+round the transport plan obtained from the Sinkhorn algorithm to a feasible transport plan in the set U(r, c), where r
+and c are the row and column marginals, respectively. The algorithm is used in the mdot.py file to round the transport
+plan and compute the cost of the rounded plan. The implementation is based on the original paper.
+"""
+
+import torch as th
+
+
+def round_altschuler(P, r, c):
+    """
+    Performs rounding given a transport plan and marginals.
+    :param P: the input transport plan
+    :param r: row marginal
+    :param c: column marginal
+    :return: rounded transport plan in feasible set U(r, c).
+    """
+    X = th.min(r / P.sum(-1), th.ones_like(r))
+    P *= X.unsqueeze(-1)
+
+    Y = th.min(c / P.sum(-2), th.ones_like(c))
+    P *= Y.unsqueeze(-2)
+
+    err_r = (r - P.sum(-1)).clamp(min=0)
+    err_c = (c - P.sum(-2)).clamp(min=0)
+    P += err_r.unsqueeze(-1) @ err_c.unsqueeze(-2) / (err_r.norm(p=1, dim=-1, keepdim=True) + 1e-30).unsqueeze(-1)
+
+    return P
+
+
+def rounded_cost_altschuler(u, v, r, c, C, gamma):
+    """Performs rounding and cost computation in logdomain given dual variables, without storing n^2 matrices.
+    :param u: dual variable for rows
+    :param v: dual variable for columns
+    :param r: row marginal
+    :param c: column marginal
+    :param C: cost matrix
+    :param gamma: temperature, i.e., the inverse of the entropic regularization weight.
+    """
+    r_P_log = u + th.logsumexp(v.unsqueeze(-2) - gamma * C, dim=-1)
+    delta_u = th.min(r.log() - r_P_log, th.zeros_like(r))
+    u += delta_u
+
+    c_P_log = v + th.logsumexp(u.unsqueeze(-1) - gamma * C, dim=-2)
+    delta_v = th.min(c.log() - c_P_log, th.zeros_like(c))
+    v += delta_v
+
+    r_P_log = u + th.logsumexp(v.unsqueeze(-2) - gamma * C, dim=-1)
+    r_P = r_P_log.exp()
+    err_r = r - r_P
+    err_r /= (err_r.norm(p=1, dim=-1, keepdim=True) + 1e-30)
+
+    c_P_log = v + th.logsumexp(u.unsqueeze(-1) - gamma * C, dim=-2)
+    c_P = c_P_log.exp()
+    err_c = c - c_P
+
+    cost = th.logsumexp(u.unsqueeze(-1) + v.unsqueeze(-2) - gamma * C + C.log(), dim=(-1, -2)).exp()
+    cost += (err_r.unsqueeze(-2) @ C @ err_c.unsqueeze(-1)).sum(-1).sum(-1)
+
+    return cost

mdot_tnt/truncated_newton.py

@@ -0,0 +1,283 @@
+
+import torch as th
+import warnings
+
+
+class TruncatedNewtonProjector:
+    def __init__(self, device, dtype, **kwargs):
+        self.device = device
+        self.rho = th.zeros(1, device=device, dtype=dtype)
+        self.debug = kwargs.get('debug', False)
+
+    def project(self, gamma_C, log_r, log_c, eps_d, u, v):
+        """
+        Project onto the set of couplings that satisfy the marginal constraints.
+        :param gamma_C: The cost matrix scaled by gamma.
+        :param log_r:
+        """
+        logs = {
+            "errs": [],
+            'ls_func_cnt': 0,
+            'chisinkhorn_steps': 0,
+            'newtonsolve_steps': 0,
+            "deltas": [],  # Ratios of actual to theoretically predicted (ideal) reduction in gradient norm.
+            "all_newtonsolve_steps": []
+        }
+        # In case of errors or issues, 10 times the tolerance level is considered
+        # a good enough solution to keep MDOT going.
+        success_fn = lambda err_: err_ < 10 * eps_d
+
+        # Each LSE operation costs 4 * n^2 operations.
+        self.LSE_r = lambda v_: th.logsumexp(v_.unsqueeze(-2) - gamma_C, dim=-1)
+        self.LSE_c = lambda u_: th.logsumexp(u_.unsqueeze(-1) - gamma_C, dim=-2)
+
+        r = log_r.exp()
+        c = log_c.exp()
+
+        log_c_P = v + self.LSE_c(u)
+        v += log_c - log_c_P  # Ensure c=c(P)
+        log_r_P = u + self.LSE_r(v)
+        k = 8
+
+        u, v, log_r_P, err, k_ = self.chi_sinkhorn(u, v, log_r, log_c, log_r_P, eps_d ** (2 / 5))
+        r_P = log_r_P.exp()
+        logs["errs"].append(err)
+        logs["chisinkhorn_steps"] = k_
+        k += k_
+
+        num_iter = 0
+
+        while err > eps_d:
+            num_iter += 1
+
+            beta = 0.5
+            eta_k = th.max(err, 0.9 * (eps_d / err))
+
+            grad_k = r_P - r
+            self.rho = th.max(th.zeros_like(self.rho), self.rho)
+
+            P = th.exp(u.unsqueeze(-1) + v.unsqueeze(-2) - gamma_C)
+            diag_PPc = ((P ** 2) / c.unsqueeze(-2)).sum(-1)
+            k += 8
+            delta_u, delta_v, matmul_cnt, rho, pcg_success = self.newton_solve(
+                P, c, diag_PPc, grad_k, r_P, err, beta, eta_k, maxIter=5000)
+            del P  # Free up memory
+            if not pcg_success:
+                k += matmul_cnt
+                logs["n_iter"] = k
+                msg = "PCG did not converge. TruncatedNewton returning with success={}".format(success_fn(err))
+                warnings.warn(msg)
+                return u, v, logs, success_fn(err)
+
+            self.rho = th.max(th.zeros_like(self.rho), 1. - (1. - rho) * 4.)
+            k += matmul_cnt
+            logs["newtonsolve_steps"] += matmul_cnt
+
+            alpha = th.ones_like(self.rho)
+            log_c_P = v + alpha * delta_v + self.LSE_c(u + alpha * delta_u)
+            k += 4
+            linear_decr = -(grad_k * delta_u).sum(-1, keepdim=True)
+            if not linear_decr > 0:
+                logs["n_iter"] = k
+                msg = "Linear decrease condition not satisfied. TruncatedNewton returning with success={}".format(
+                    success_fn(err))
+                warnings.warn(msg)
+                return u, v, logs, success_fn(err)
+
+            armijo = log_c_P.exp().sum(-1, keepdim=True) - 1 <= 0.99 * alpha * linear_decr
+            while not armijo:  # Check armijo condition for batch elements where err > eps_d
+                alpha *= 0.5
+                if alpha < 1e-9:
+                    logs["n_iter"] = k
+                    msg = "Line search did not converge. TruncatedNewton returning with success={}".format(
+                        success_fn(err))
+                    warnings.warn(msg)
+                    return u, v, logs, success_fn(err)
+
+                log_c_P = v + alpha * delta_v + self.LSE_c(u + alpha * delta_u)
+                k += 4
+                logs["ls_func_cnt"] += 4
+                armijo = log_c_P.exp().sum(-1, keepdim=True) - 1 <= 0.99 * alpha * linear_decr
+
+            u += alpha * delta_u
+            v += alpha * delta_v
+
+            # The below error (before the Sinkhorn update) is used
+            # to measure the progress of the algorithm with TruncatedNewton steps.
+            err_before_sk = (c - log_c_P.exp()).abs().sum(-1)
+            err_before_sk += (r - (u + self.LSE_r(v)).exp()).abs().sum(-1)
+
+            # Sinkhorn update to ensure c=c(P).
+            v += log_c - log_c_P
+
+            log_r_P = u + self.LSE_r(v)
+            k += 4
+
+            u, v, log_r_P, err, k_ = self.chi_sinkhorn(u, v, log_r, log_c, log_r_P, eps_d ** (2 / 5))
+            r_P = log_r_P.exp()
+            logs["chisinkhorn_steps"] += k_
+            k += k_
+
+            logs["errs"].append(err)
+            logs["deltas"].append(th.min((logs["errs"][-2] - err_before_sk) / ((1 - eta_k) * logs["errs"][-2])).item())
+
+        if u.isnan().any() or v.isnan().any():
+            raise ValueError("NaNs encountered in u or v")
+
+        logs["n_iter"] = k
+
+        # Since we already computed log_r_P, we can use it to perform one last Sinkhorn update on rows.
+        delta_u = log_r - log_r_P
+        u += delta_u
+
+        return u, v, logs, True
+
+    def chi_sinkhorn(self, u, v, log_r, log_c, log_r_P, eps_chi, maxOps=float('inf')):
+        k = 0
+        r = log_r.exp()
+        err = (r - log_r_P.exp()).norm(p=1, dim=-1)
+        r_P = log_r_P.exp()
+        chi_squared = ((r - r_P) ** 2 / r_P).sum(-1)
+
+        while chi_squared > eps_chi and k < maxOps:
+            u += log_r - log_r_P
+
+            log_c_P = v + self.LSE_c(u)
+            v += log_c - log_c_P
+
+            log_r_P = u + self.LSE_r(v)
+            r_P = log_r_P.exp()
+
+            err = (r - r_P).norm(p=1, dim=-1)
+
+            chi_squared = ((r - r_P) ** 2 / r_P).sum(-1)
+            k += 8
+
+        if k >= maxOps:
+            raise ValueError("Chi-Sinkhorn did not converge in maxIter={} steps".format(maxOps))
+
+        return u, v, log_r_P, err, k
+
+    def newton_solve(self, P, c, diag_PPc, grad_k, r_P, err, beta=0.5, eta_k=0.5, maxIter=500):
+        rho = self.rho
+        tol = err * eta_k
+
+        def matmul_PPc(x_):
+            return (P @ ((x_.unsqueeze(-2) @ P).transpose(-2, -1) / c.unsqueeze(-1))).squeeze(-1)
+
+        # mml = th.compile(matmul_PPc)
+        mml = matmul_PPc
+
+        M = lambda rho_: (r_P - rho_ * diag_PPc)  # Diagonal preconditioner
+        M_rho = M(th.ones_like(self.rho))
+        M_rho[M_rho <= 0] = M_rho[M_rho > 0].min()
+
+        x0 = -grad_k / M_rho
+        PPc_x0 = mml(x0)
+        matmul_cnt = 2
+        r_P_x0 = (r_P * x0)
+
+        x = x0.clone()
+        PPc_x = PPc_x0.clone()
+        r_P_x = r_P_x0.clone()
+
+        res_true = r_P_x0 - PPc_x + grad_k
+
+        linear_decr = (x * -grad_k).sum(-1)
+        if linear_decr <= 0:
+            raise ValueError("Linear decrease condition not satisfied")
+
+        r_true_norm = res_true.norm(p=1, dim=-1)
+        best_sol = x.clone()
+        best_r_true_norm = r_true_norm.clone()
+
+        done = False
+        success = True
+
+        while best_r_true_norm > tol:
+            best_sol[r_true_norm < best_r_true_norm] = x[r_true_norm < best_r_true_norm]
+            best_r_true_norm = th.min(r_true_norm, best_r_true_norm)
+
+            rho[r_true_norm > tol] = 1. - (1. - rho[r_true_norm > tol]) * 0.25
+            M_rho = M(rho)
+
+            if matmul_cnt > 0:
+                x = x0.clone()
+                PPc_x = PPc_x0.clone()
+                r_P_x = r_P_x0.clone()
+
+            Fr_x = r_P_x - rho * PPc_x
+            res = Fr_x + grad_k
+
+            res_true = r_P_x - PPc_x + grad_k
+            r_true_norm = res_true.norm(p=1, dim=-1)
+            best_r_true_norm = th.min(r_true_norm, best_r_true_norm)
+            linear_decr = (x * -grad_k).sum(-1)
+            if (best_r_true_norm < tol).all() and (linear_decr > 0).all():
+                break
+
+            y = res / M_rho
+            p = -y.clone()
+            ry_old = (res * y).sum(-1, keepdim=True)
+
+            r_norm = res.norm(p=1, dim=-1)
+            while (r_norm > 0.5 * (1 - beta) * tol)[best_r_true_norm > tol].any():
+                PPc_p = mml(p)
+                matmul_cnt += 2
+                Fr_p = (r_P * p) - rho * PPc_p
+
+                quad = (Fr_p * p).sum(-1, keepdim=True)
+                if (quad <= 0)[best_r_true_norm > tol].any():
+                    warnings.warn("Warning: negative curvature encountered in CG. Returning best solution. "
+                          "Residual norm less than error: {}".format((best_r_true_norm < err).item()))
+                    x = best_sol.clone()
+                    done = True
+                    success = best_r_true_norm < err
+                    warnings.warn("Resetting discount factor rho = 0")
+                    rho = th.zeros_like(self.rho)
+                    break
+
+                alpha = ry_old / quad
+                x += alpha * p
+                res += alpha * Fr_p
+                r_norm = res.norm(p=1, dim=-1)
+
+                if th.isnan(r_norm)[best_r_true_norm > tol].any() or th.isinf(r_norm)[best_r_true_norm > tol].any():
+                    raise ValueError("NaNs or infs encountered in r_norm")
+
+                PPc_x += alpha * PPc_p
+
+                r_P_x = (r_P * x)
+                res_true = r_P_x - PPc_x + grad_k
+                r_true_norm = res_true.norm(p=1, dim=-1)
+                best_sol[r_true_norm < best_r_true_norm] = x[r_true_norm < best_r_true_norm]
+                best_r_true_norm = th.min(r_true_norm, best_r_true_norm)
+
+                linear_decr = (x * -grad_k).sum(-1)
+                if (best_r_true_norm <= tol).all() and (linear_decr > 0).all():
+                    done = True
+                    success = True
+                    break
+
+                if matmul_cnt > 2 * maxIter:
+                    warnings.warn("PCG did not converge.")
+                    done = True
+                    success = False
+                    break
+
+                y = res / M_rho
+                ry_new = (res * y).sum(-1, keepdim=True)
+                p = -y + (ry_new / ry_old) * p
+                ry_old = ry_new.clone()
+
+            if done:
+                break
+
+        if r_true_norm <= tol:
+            success = True
+
+        x = best_sol
+        Pc_x = ((x.unsqueeze(-2) @ P).transpose(-2, -1) / c.unsqueeze(-1)).squeeze(-1)
+        matmul_cnt += 1
+
+        return x, -Pc_x, matmul_cnt, rho, success

mdot_tnt-0.1.0.dist-info/LICENSE

@@ -0,0 +1,26 @@
+# Non-Commercial Research License (NCRL-1.0)
+
+Copyright (C) 2025 Mete Kemertas
+
+## 1. License Grant
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to use, copy, modify, merge, publish, and distribute copies of the Software **solely for non-commercial research, educational, and personal purposes**, subject to the following conditions:
+
+## 2. Restrictions
+### 2.1 **Non-Commercial Use Only**
+- The Software **may NOT** be used for any commercial purpose without explicit written permission from the Licensor.
+- "Commercial purpose" includes, but is not limited to:
+  - Selling or licensing the Software.
+  - Using the Software in proprietary products or services.
+  - Offering the Software as part of a paid or revenue-generating service.
+
+### 2.2 **No Warranty & Liability**
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES, OR OTHER LIABILITY ARISING FROM THE USE OF THE SOFTWARE.
+
+### 2.3 **Commercial Licensing**
+For commercial use, a separate license must be obtained from the Licensor. To inquire about licensing, please contact: **kemertas@cs.toronto.edu**.
+
+## 3. Termination
+This license automatically terminates if the Licensee breaches any of its terms. Upon termination, all rights granted under this license are revoked, and the Licensee must cease using and distributing the Software.
+
+## 4. Governing Law and Enforcement
+This license shall be governed by and construed in accordance with the laws of Ontario, Canada. However, violations of this license may also be pursued un

mdot_tnt-0.1.0.dist-info/METADATA

@@ -0,0 +1,71 @@
+Metadata-Version: 2.1
+Name: mdot-tnt
+Version: 0.1.0
+Summary: A fast, GPU-parallel, PyTorch-compatible optimal transport solver.
+Author-email: Mete Kemertas <kemertas@cs.toronto.edu>
+Project-URL: Homepage, https://github.com/metekemertas/mdot_tnt
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+License-File: LICENSE
+
+This is the official repository for the MDOT-TruncatedNewton (or MDOT-TNT)
+algorithm [1] for solving the entropic-regularized optimal transport (OT) problem. 
+In addition to being GPU-friendly, the algorithm is stable under weak regularization and can therefore find highly
+precise approximations of the un-regularized problem's solution quickly. 
+
+The current implementation is based on PyTorch and is compatible with both CPU and GPU. PyTorch is the only dependency.
+
+
+For installation:
+First, install PyTorch following the instructions at https://pytorch.org/get-started/locally/ to select the version that matches your system's configuration.
+```bash
+pip install mdot_tnt
+```
+
+Quickstart guide:
+```
+import mdot_tnt
+import torch as th
+device = 'cuda' if th.cuda.is_available() else 'cpu'
+N, M, dim = 100, 200, 128
+
+# Sample row and column marginals from Dirichlet distributions
+r = th.distributions.Dirichlet(th.ones(N)).sample()
+c = th.distributions.Dirichlet(th.ones(M)).sample()
+
+# Cost matrix from pairwise Euclidean distances of random points in R^100
+x = th.distributions.MultivariateNormal(th.zeros(dim), th.eye(dim)).sample((N,))
+y = th.distributions.MultivariateNormal(th.zeros(dim), th.eye(dim)).sample((M,))
+C = th.cdist(x, y, p=2)
+C /= C.max()  # Normalize cost matrix to meet convention.
+
+# Use double precision for numerical stability in high precision regime.
+r, c, C = r.double(), c.double(), C.double()
+
+# Solve OT problem. Increase (decrease) gamma_f for higher (lower) precision.
+# Default is gamma_f=2**12. Expect error of order logn / gamma_f at worst, and possibly lower.
+cost = mdot_tnt.solve_OT(r, c, C, gamma_f=2**12)
+
+# To return a feasible transport plan, use the following:
+transport_plan = mdot_tnt.solve_OT(r, c, C, gamma_f=2**12, return_plan=True)
+
+# In both cases, the default rounding onto the feasible set can be disabled by setting `round=False`.
+```
+
+The code is released under a custom non-commerical use license. If you use our work in
+your research, please consider citing:
+
+```
+@inproceedings{
+kemertas2025a,
+title={A Truncated Newton Method for Optimal Transport},
+author={Mete Kemertas and Amir-massoud Farahmand and Allan Douglas Jepson},
+booktitle={The Thirteenth International Conference on Learning Representations},
+year={2025},
+url={https://openreview.net/forum?id=gWrWUaCbMa}
+}
+```
+
+For inquiries, email: kemertas [at] cs [dot] toronto [dot] edu
+
+[1] Mete Kemertas, Amir-massoud Farahmand, Allan Douglas Jepson. "A Truncated Newton Method for Optimal Transport." The Thirteenth International Conference on Learning Representations (ICLR), 2025. https://openreview.net/forum?id=gWrWUaCbMa

mdot_tnt-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+mdot_tnt/__init__.py,sha256=zi0ujbN19nv_dt8wdxPQgsOuDi-4gIg4sq6y9fDfAkI,2657
+mdot_tnt/mdot.py,sha256=EUch7ECZxDWDoNW3S84qyCx6-qbG6Mu1_XlWow862Wc,4620
+mdot_tnt/rounding.py,sha256=Q7QBPsFzBqnMZKnlV147ruStUCme6gwt6HnjTjqBezk,2405
+mdot_tnt/truncated_newton.py,sha256=YEkWJxL5H4XsuWQzfdZ0E5Ju3Npk3kKoe2l_go8wq2A,10314
+mdot_tnt-0.1.0.dist-info/LICENSE,sha256=sXw3FpVqouAddNhfwD6nXaSgABFyLnXuSn_ghjU0AhY,1837
+mdot_tnt-0.1.0.dist-info/METADATA,sha256=2Ou4X10K6bUGHQhnMXoFjrv-4YUeIyreSKibMLM2gq0,2972
+mdot_tnt-0.1.0.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+mdot_tnt-0.1.0.dist-info/top_level.txt,sha256=HmxTNtoLH7F20hgZVFdfUowIQ2fviSX64wSG1HP8Iao,9
+mdot_tnt-0.1.0.dist-info/RECORD,,

mdot_tnt-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.3.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

mdot_tnt-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+mdot_tnt