polystep 0.7.0__py3-none-any.whl

polystep/solvers/sinkhorn.py

@@ -0,0 +1,522 @@
+"""Log-domain Sinkhorn solver for entropic optimal transport.
+
+PolyStep's OT problems are (n particles x m=V polytope vertices) with V small
+(2*particle_dim), so the cost is O(n*V) -- a full-rank log-domain solve is
+already optimal and there is no large dense kernel to approximate away.
+"""
+
+import functools
+import math
+import warnings
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import torch
+
+from ..costs import scale_cost_matrix
+from ._prelude import align_dual, align_marginal, recenter_cost, sanitize_cost, validate_positive
+
+
+@dataclass
+class SinkhornResult:
+    """Output from a Sinkhorn solve.
+
+    Attributes:
+        f: First dual potential of shape (n,). Together with ``g``, these
+            dual potentials encode the optimal transport solution. They can
+            be reused as warm-start initializations for the next solve,
+            which typically reduces iterations from ~100 to ~10.
+        g: Second dual potential of shape (m,). See ``f``.
+        converged: Whether the solver converged within tolerance.
+        n_iters: Number of iterations actually run.
+        ent_reg_cost: Entropic dual objective <f, a> + <g, b> - eps*sum_ij P_ij
+            (true plan mass), reported in the original cost frame. Equals the
+            regularized transport cost at convergence.
+        errors: Per-check marginal errors (for diagnostics).
+    """
+
+    f: torch.Tensor
+    g: torch.Tensor
+    converged: bool
+    n_iters: int
+    ent_reg_cost: float
+    errors: Optional[List[float]] = None
+
+    # Internal fields for lazy .matrix computation
+    _eps: float = float("nan")
+    _cost_matrix: Optional[torch.Tensor] = None  # (n, m)
+
+    @functools.cached_property
+    def matrix(self) -> torch.Tensor:
+        """Transport plan, computed lazily: P_ij = exp((f_i + g_j - C_ij) / eps)."""
+        log_P = (self.f.unsqueeze(1) + self.g.unsqueeze(0) - self._cost_matrix) / self._eps
+        return torch.exp(log_P)
+
+
+@dataclass
+class SinkhornSolver:
+    """Full-rank log-domain Sinkhorn solver for entropic optimal transport.
+
+    Solves the entropic optimal transport problem by alternating row and
+    column scaling in log domain. The entropic regularization parameter
+    ``epsilon`` controls the trade-off between transport cost minimization
+    and entropy maximization: higher epsilon gives a smoother, more diffuse
+    transport plan (easier to solve but less precise), while lower epsilon
+    gives a sharper plan closer to exact OT (but harder to solve numerically).
+
+    The solver converges when the marginal constraint violation falls below
+    ``threshold``. Warm-starting with previous dual potentials (f, g) from
+    ``SinkhornResult`` typically reduces iterations from ~100 to ~10.
+
+    Attributes:
+        epsilon: Entropic regularization strength.
+        max_iterations: Maximum number of Sinkhorn iterations.
+        threshold: Convergence threshold on marginal error.
+            Set <= 0 for fixed-iteration mode (no early stopping).
+        check_every: Check convergence every N iterations.
+        compile: Whether to use torch.compile for hot paths (requires CUDA).
+    """
+
+    epsilon: float = 0.1
+    max_iterations: int = 2000
+    threshold: float = 1e-6
+    check_every: int = 10
+    compile: bool = False
+    omega: float = 1.0
+    anderson_depth: int = 0  # 0 = disabled, >0 = ring buffer depth for Anderson acceleration
+    adaptive_omega: bool = False  # False = static omega, True = residual-ratio dynamic omega (Lehmann 2022)
+    data_dependent_init: bool = False  # False = zeros init, True = cost-mean init for cold starts
+
+    def __post_init__(self):
+        """Initialize compiled function registry and validate parameters."""
+        if self.epsilon <= 0:
+            raise ValueError(
+                f"epsilon must be > 0, got {self.epsilon}. "
+                f"A zero or negative epsilon causes division by zero in log-domain Sinkhorn iterations."
+            )
+        if self.omega < 0.5 or self.omega > 1.95:
+            raise ValueError(
+                f"omega must be in [0.5, 1.95], got {self.omega}. "
+                f"Values < 0.5 cause divergence; values > 1.95 are numerically unstable. "
+                f"Recommended range: [1.0, 1.8] for acceleration."
+            )
+        if self.check_every < 1:
+            raise ValueError(
+                f"check_every must be >= 1, got {self.check_every}. It is the modulus for periodic convergence checks."
+            )
+        if self.max_iterations < 1:
+            raise ValueError(
+                f"max_iterations must be >= 1, got {self.max_iterations}. "
+                f"Zero iterations leave the duals at their init and return an infeasible plan."
+            )
+
+        from .._compiled import CompiledFunctions
+
+        self._compiled = CompiledFunctions(compile=self.compile and torch.cuda.is_available())
+
+    def solve(
+        self,
+        cost_matrix: torch.Tensor,
+        a: Optional[torch.Tensor] = None,
+        b: Optional[torch.Tensor] = None,
+        init_f: Optional[torch.Tensor] = None,
+        init_g: Optional[torch.Tensor] = None,
+        scale_cost: Optional[Union[str, float]] = None,
+        init_eps: Optional[float] = None,
+    ) -> SinkhornResult:
+        """Solve entropic OT problem.
+
+        Solves: min_P <C, P> + eps * KL(P || a x b)
+        s.t.  P 1 = a,  P^T 1 = b,  P >= 0
+
+        Args:
+            cost_matrix: Cost matrix C of shape (n, m).
+            a: Source marginal of shape (n,). Defaults to uniform 1/n.
+            b: Target marginal of shape (m,). Defaults to uniform 1/m.
+            init_f: Warm-start first dual potential of shape (n,).
+            init_g: Warm-start second dual potential of shape (m,).
+            scale_cost: Cost scaling: 'mean', 'max_cost', float, or None.
+            init_eps: Epsilon under which ``init_f`` / ``init_g`` were
+                computed by a previous solve. When provided and different
+                from ``self.epsilon``, the duals are rescaled by
+                ``self.epsilon / init_eps`` (holding the log-domain scaling
+                ``u = f / eps`` fixed). This is a numerical *heuristic*, not
+                an exact warm start -- entropic cost-unit potentials are not
+                homogeneous in epsilon. Benchmarks
+                (``experiments/scripts/bench_eps_rescale.py``) show it is
+                ~neutral for well-scaled costs and modestly reduces
+                non-converged solves in the large-cost / small-epsilon
+                regime, where un-rescaled duals risk ``exp`` overflow.
+
+        Returns:
+            SinkhornResult with dual potentials and transport plan access.
+        """
+        # Schedules mutate ``self.epsilon`` per step, so re-validate here (a
+        # plain float check, no host sync) rather than only at construction.
+        validate_positive(
+            self.epsilon,
+            "epsilon",
+            "A zero/negative epsilon divides by zero in log-domain Sinkhorn.",
+        )
+        return self._solve_full_rank(
+            cost_matrix,
+            a,
+            b,
+            init_f,
+            init_g,
+            scale_cost,
+            init_eps,
+        )
+
+    def _solve_full_rank(
+        self,
+        cost_matrix: torch.Tensor,
+        a: Optional[torch.Tensor],
+        b: Optional[torch.Tensor],
+        init_f: Optional[torch.Tensor],
+        init_g: Optional[torch.Tensor],
+        scale_cost: Optional[Union[str, float]],
+        init_eps: Optional[float] = None,
+    ) -> SinkhornResult:
+        """Full-rank log-domain Sinkhorn iterations."""
+        # Shared prelude: FP32 promotion (BF16's 7 mantissa bits collapse the
+        # log-sum-exp row-max trick past ~15 nats) + device-side non-finite
+        # cost handling (no host sync), then device/dtype-aligned marginals.
+        cost_matrix = sanitize_cost(cost_matrix)
+        n, m = cost_matrix.shape
+        if n == 0 or m == 0:
+            raise ValueError(
+                f"Sinkhorn received an empty cost matrix (shape {(n, m)}); "
+                "at least one source and target point are required."
+            )
+        device, dtype = cost_matrix.device, cost_matrix.dtype
+        # Balanced Sinkhorn needs sum(a) == sum(b). Only check when the caller
+        # gave both marginals; the default uniform a=b=None is balanced, so the
+        # hot path stays free of a host sync.
+        marginals_given = a is not None and b is not None
+        a = align_marginal(a, n, device, dtype, "a")
+        b = align_marginal(b, m, device, dtype, "b")
+        if marginals_given and not torch.isclose(a.sum(), b.sum(), rtol=1e-3, atol=1e-6):
+            warnings.warn(
+                f"Sinkhorn marginals have unequal total mass (a.sum()={a.sum().item():.6g}, "
+                f"b.sum()={b.sum().item():.6g}); balanced OT is infeasible, so 'converged' "
+                f"and the returned plan may be meaningless. Use KLSoftmaxSolver for unbalanced OT.",
+                stacklevel=2,
+            )
+
+        # Scale cost matrix (division creates a new tensor; no clone needed)
+        cost_matrix = scale_cost_matrix(cost_matrix, scale_cost)
+        # Recenter to min(C)=0 so the log-sum-exp and the exp((f+g-C)/eps) plan
+        # keep FP32 precision when |C| is much larger than eps. cost_shift is
+        # added back into ent_reg_cost below.
+        cost_matrix, cost_shift = recenter_cost(cost_matrix)
+
+        # Log kernel: log_K = -C / eps
+        eps = self.epsilon
+        log_K = -cost_matrix / eps
+
+        log_a = torch.log(torch.clamp(a, min=1e-30))
+        log_b = torch.log(torch.clamp(b, min=1e-30))
+
+        # Initialize dual potentials with warm-start validation
+        # Data-dependent initialization: set initial duals from cost matrix means
+        # Only when no warm-start provided. Applied AFTER cost scaling on the scaled matrix.
+        if self.data_dependent_init and init_f is None and init_g is None:
+            f = -cost_matrix.mean(dim=1)
+            g = -cost_matrix.mean(dim=0)
+        else:
+            # align_dual moves the warm start onto (device, dtype) and clones
+            # it; it returns None on a shape mismatch -> fall back to zeros.
+            f = align_dual(init_f, n, device, dtype, "init_f")
+            g = align_dual(init_g, m, device, dtype, "init_g")
+            if f is None:
+                f = torch.zeros(n, device=device, dtype=dtype)
+            if g is None:
+                g = torch.zeros(m, device=device, dtype=dtype)
+
+        # Rescale the warm-started duals when the caller changed epsilon since
+        # the previous solve (see ``init_eps`` docstring above). Do this before
+        # the clamp below so a large ``eps/init_eps`` ratio stays bounded and
+        # cannot blow the plan to Inf/NaN.
+        if init_eps is not None and init_eps > 0 and (init_f is not None or init_g is not None):
+            if abs(init_eps - eps) / max(eps, 1e-9) > 1e-6:
+                scale_factor = eps / init_eps
+                f = f * scale_factor
+                g = g * scale_factor
+
+        # Validate warm-started dual potentials. Dual potentials scale with the
+        # cost matrix magnitude, not epsilon. ``cost_scale`` is kept on-device so
+        # the clamp does not sync every solve. Runs after the rescale above so the
+        # bound holds.
+        cost_scale = cost_matrix.abs().max().clamp(min=1e-6)
+        max_abs_dual = 10.0 * cost_scale
+        if not (torch.isfinite(f).all() and torch.isfinite(g).all()):
+            f.zero_()
+            g.zero_()
+        else:
+            f.clamp_(-max_abs_dual, max_abs_dual)
+            g.clamp_(-max_abs_dual, max_abs_dual)
+
+        # Re-center using the only valid dual gauge ``f -> f + c, g -> g - c``,
+        # which leaves ``f_i + g_j`` (and hence the plan ``P``) unchanged. The
+        # half-difference shift balances the two potentials' magnitudes so the
+        # warm-start clamp above keeps them bounded across a long schedule.
+        # NB: independent mean subtraction (``f -= f.mean(); g -= g.mean()``)
+        # is NOT a valid gauge and perturbs the iterate under overrelaxation.
+        c = 0.5 * (g.mean() - f.mean())
+        f = f + c
+        g = g - c
+
+        # Determine if we should check convergence
+        fixed_mode = self.threshold <= 0
+
+        if fixed_mode:
+            if self.anderson_depth > 0:
+                warnings.warn(
+                    "anderson_depth > 0 has no effect in fixed-iteration mode "
+                    "(threshold <= 0). Anderson acceleration is only supported "
+                    "in convergence-checking mode.",
+                    stacklevel=2,
+                )
+            if self.adaptive_omega:
+                warnings.warn(
+                    "adaptive_omega=True has no effect in fixed-iteration mode "
+                    "(threshold <= 0). Adaptive omega is only supported "
+                    "in convergence-checking mode.",
+                    stacklevel=2,
+                )
+
+        converged = False
+        n_iters = 0
+        errors: List[float] = []
+
+        omega = self.omega
+
+        # Pin the iteration loop inside an autocast-disabled FP32
+        # region so a caller running under ``autocast(bfloat16)`` can't
+        # demote our log-sum-exp intermediates.
+        with torch.no_grad(), torch.amp.autocast("cuda", enabled=False), torch.amp.autocast("cpu", enabled=False):
+            if fixed_mode:
+                # Fixed-iteration path: compiled body with post-loop NaN check.
+                # Warm-started solver with well-conditioned cost matrix rarely
+                # diverges mid-iteration; checking only at the end avoids
+                # GPU-CPU sync overhead from periodic isfinite() reductions.
+                sinkhorn_iter = self._compiled.sinkhorn_iter
+                for i in range(self.max_iterations):
+                    f, g = sinkhorn_iter(f, g, log_K, log_a, log_b, eps, omega)
+                    n_iters = i + 1
+                # Fixed mode: a finite result is success. converged=False here
+                # would make ProgressiveEpsilon inflate epsilon.
+                if not (torch.isfinite(f).all() and torch.isfinite(g).all()):
+                    f.zero_()
+                    g.zero_()
+                else:
+                    converged = True
+            else:
+                # Convergence-checking path: stays fully eager with overrelaxation
+                # Anderson acceleration: ring buffer for iterate mixing
+                if self.anderson_depth > 0:
+                    aa_history_x = []  # list of (f, g) pairs
+                    aa_history_r = []  # list of (r_f, r_g) residual pairs
+
+                # Adaptive omega: residual-ratio estimator state (Lehmann 2022)
+                if self.adaptive_omega:
+                    prev_err = None
+
+                # Divergence detector for static omega. Track consecutive
+                # growths of ``|f|.max + |g|.max``; if the iterate norm
+                # keeps growing across ``_divergence_patience`` checks,
+                # back omega off to 1.0 (Lehmann 2022's proven-safe value).
+                _divergence_prev_norm = float("inf")
+                _divergence_growth_count = 0
+                _divergence_patience = 3
+                # Latch the divergence back-off so the adaptive estimator below
+                # cannot re-raise omega in the same check.
+                _omega_capped = False
+
+                def dual_objective(fv, gv):
+                    # D(f,g) = <f,a> + <g,b> - eps * sum_ij exp((f_i+g_j-C_ij)/eps).
+                    # The true Sinkhorn Lyapunov: unlike <f,a>+<g,b> alone, it is
+                    # valid off the marginal constraint (sum(P) != 1).
+                    log_P = fv.unsqueeze(1) / eps + gv.unsqueeze(0) / eps + log_K
+                    mass_p = torch.exp(torch.logsumexp(log_P.reshape(-1), dim=0))
+                    return (fv * a).sum() + (gv * b).sum() - eps * mass_p
+
+                for i in range(self.max_iterations):
+                    f_target = eps * (log_a - torch.logsumexp(log_K + g.unsqueeze(0) / eps, dim=1))
+                    f_new = (1 - omega) * f + omega * f_target
+                    g_target = eps * (log_b - torch.logsumexp(log_K + f_new.unsqueeze(1) / eps, dim=0))
+                    g_new = (1 - omega) * g + omega * g_target
+
+                    # Anderson acceleration
+                    if self.anderson_depth > 0 and (i + 1) % self.check_every == 0:
+                        r_f = f_new - f
+                        r_g = g_new - g
+                        aa_history_x.append((f.clone(), g.clone()))
+                        aa_history_r.append((r_f.clone(), r_g.clone()))
+                        m_depth = self.anderson_depth
+                        if len(aa_history_x) > m_depth + 1:
+                            aa_history_x.pop(0)
+                            aa_history_r.pop(0)
+
+                        if len(aa_history_r) >= 2:
+                            k = len(aa_history_r) - 1
+                            # Build residual difference matrix
+                            delta_r = torch.stack(
+                                [
+                                    torch.cat(
+                                        [
+                                            aa_history_r[j + 1][0] - aa_history_r[j][0],
+                                            aa_history_r[j + 1][1] - aa_history_r[j][1],
+                                        ]
+                                    )
+                                    for j in range(k)
+                                ],
+                                dim=1,
+                            )  # (n+m, k)
+                            current_r = torch.cat([r_f, r_g])  # (n+m,)
+
+                            # Tikhonov-regularized least squares via the augmented
+                            # system [dR; sqrt(lambda) I] alpha = [r; 0], solved by
+                            # QR (more stable than the normal equations). lambda
+                            # scales with ||dR||^2 to track conditioning.
+                            try:
+                                lam = 1e-8 * (delta_r * delta_r).sum().clamp(min=1e-30)
+                                aug_A = torch.cat(
+                                    [
+                                        delta_r,
+                                        torch.sqrt(lam) * torch.eye(k, device=delta_r.device, dtype=delta_r.dtype),
+                                    ],
+                                    dim=0,
+                                )
+                                aug_b = torch.cat(
+                                    [current_r, torch.zeros(k, device=delta_r.device, dtype=delta_r.dtype)]
+                                )
+                                alpha = torch.linalg.lstsq(aug_A, aug_b.unsqueeze(1)).solution.squeeze(1)  # (k,)
+
+                                # Guard against NaN/Inf and huge alpha from ill-conditioning
+                                if torch.isfinite(alpha).all() and alpha.norm() < 1e3:
+                                    delta_x = torch.stack(
+                                        [
+                                            torch.cat(
+                                                [
+                                                    aa_history_x[j + 1][0] - aa_history_x[j][0],
+                                                    aa_history_x[j + 1][1] - aa_history_x[j][1],
+                                                ]
+                                            )
+                                            for j in range(k)
+                                        ],
+                                        dim=1,
+                                    )
+                                    combined = torch.cat([f_new, g_new]) - delta_x @ alpha
+                                    # Validate combined result before assigning
+                                    if torch.isfinite(combined).all():
+                                        # Accept the Anderson step only when its dual
+                                        # objective beats the plain iterate and does
+                                        # not regress below the previous iterate. The
+                                        # previous-iterate check stays monotone under
+                                        # overrelaxation, where the plain SOR step is
+                                        # not. Device-side compare avoids a sync.
+                                        f_combined = combined[:n]
+                                        g_combined = combined[n:]
+                                        lyap_prev = dual_objective(f, g)
+                                        lyap_plain = dual_objective(f_new, g_new)
+                                        lyap_combined = dual_objective(f_combined, g_combined)
+                                        accept = (lyap_combined >= lyap_plain - 1e-6) & (
+                                            lyap_combined >= lyap_prev - 1e-6
+                                        )
+                                        f_new = torch.where(accept, f_combined, f_new)
+                                        g_new = torch.where(accept, g_combined, g_new)
+                            except RuntimeError:
+                                pass  # Fall back to standard iterate on solver failure (expected for ill-conditioned problems)
+
+                    f, g = f_new, g_new
+                    n_iters = i + 1
+
+                    # Convergence and divergence check (every check_every iterations
+                    # to avoid GPU-CPU sync on every iteration). All scalar
+                    # measurements are batched into a single device->host
+                    # transfer per check to amortize the sync cost.
+                    if (i + 1) % self.check_every == 0:
+                        # Divergence check
+                        if torch.isnan(f).any() or torch.isinf(f).any() or torch.isnan(g).any() or torch.isinf(g).any():
+                            f.zero_()
+                            g.zero_()
+                            break
+
+                        log_P_row = f.unsqueeze(1) / eps + log_K + g.unsqueeze(0) / eps
+                        marginal_a = torch.exp(torch.logsumexp(log_P_row, dim=1))
+                        marginal_b = torch.exp(torch.logsumexp(log_P_row, dim=0))
+
+                        # Batch all scalar measurements into one transfer.
+                        err_a_t = torch.max(torch.abs(marginal_a - a))
+                        err_b_t = torch.max(torch.abs(marginal_b - b))
+                        if omega > 1.5:
+                            dual_norm_t = f.abs().max() + g.abs().max()
+                        else:
+                            dual_norm_t = err_a_t  # placeholder, value unused
+                        err_a, err_b, dual_norm_v = torch.stack([err_a_t, err_b_t, dual_norm_t]).tolist()
+                        err = max(err_a, err_b)
+
+                        # Static-omega divergence detector. Lehmann et al.
+                        # 2022 give a safe range ``omega in (0, 2 - rho)``
+                        # where rho is the linearised spectral radius.
+                        # Empirically ``omega <= 1.5`` is safe on
+                        # well-conditioned C; only monitor at the aggressive
+                        # end. Require both sustained growth (>5% per check)
+                        # and three-in-a-row to avoid firing on the benign
+                        # iterate ripples that occur near the fixed point.
+                        if omega > 1.5:
+                            if dual_norm_v > _divergence_prev_norm * 1.05:
+                                _divergence_growth_count += 1
+                                if _divergence_growth_count >= _divergence_patience:
+                                    warnings.warn(
+                                        f"Sinkhorn divergence detected with "
+                                        f"omega={omega:.2f} after "
+                                        f"{_divergence_patience} consecutive "
+                                        f"growth checks (>5% per check); "
+                                        f"backing omega off to 1.0 (safe).",
+                                        stacklevel=2,
+                                    )
+                                    omega = 1.0
+                                    _omega_capped = True
+                                    _divergence_growth_count = 0
+                            else:
+                                _divergence_growth_count = 0
+                            _divergence_prev_norm = dual_norm_v
+
+                        # Adaptive omega via the Lehmann residual-ratio estimator
+                        # (arXiv:2012.12562): estimate the linear rate from the
+                        # ratio of successive marginal errors and pick the optimal
+                        # overrelaxation. Converges to omega_opt instead of
+                        # oscillating like a fixed step-up/step-down heuristic.
+                        if self.adaptive_omega and not _omega_capped:
+                            if prev_err is not None and prev_err > 1e-12 and err > 0.0:
+                                r = min(err / prev_err, 0.99)
+                                omega = 2.0 / (1.0 + math.sqrt(1.0 - r ** (1.0 / self.check_every)))
+                                omega = min(max(omega, 1.0), 1.95)
+                            prev_err = err
+
+                        errors.append(err)
+                        if err < self.threshold:
+                            converged = True
+                            break
+
+        # Dual objective <f,a> + <g,b> - eps*sum_ij P_ij, using the true plan mass
+        # (equals eps*sum(a) only at convergence). cost_shift*sum(a) undoes the
+        # recenter so the value is in the original cost frame.
+        log_P_final = f.unsqueeze(1) / eps + log_K + g.unsqueeze(0) / eps
+        plan_mass = torch.exp(torch.logsumexp(log_P_final.reshape(-1), dim=0))
+        ent_reg_cost = ((f * a).sum() + (g * b).sum() - eps * plan_mass + cost_shift * a.sum()).item()
+
+        result = SinkhornResult(
+            f=f,
+            g=g,
+            converged=converged,
+            n_iters=n_iters,
+            ent_reg_cost=ent_reg_cost,
+            errors=errors if errors else None,
+            _eps=eps,
+            _cost_matrix=cost_matrix,
+        )
+        return result

polystep/solvers/softmax.py

@@ -0,0 +1,151 @@
+"""SoftmaxSolver: direct softmax weighting for subspace modes.
+
+Replaces iterative Sinkhorn OT with a single softmax(-C/epsilon) pass.
+Mathematically equivalent when the OT target marginal constraint is
+naturally satisfied (few particles in subspace mode). See paper Section 5.10.
+
+Key properties:
+    - Row sums of transport matrix equal source marginal a
+    - No dual potentials (f, g are None)
+    - Single iteration (converged=True, n_iters=1)
+    - Numerical stability via PyTorch's built-in softmax (subtracts row-max)
+"""
+
+import warnings
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+
+from ..costs import scale_cost_matrix
+from ._prelude import align_marginal, recenter_cost, sanitize_cost, validate_positive
+from .base import SolverResult
+
+
+# Below this ratio of `epsilon` to `max|C|`, ``-C/epsilon`` overflows
+# before ``torch.softmax`` gets a chance to subtract the row max.
+# Tested empirically at FP32 / BF16: ratios above ~1e-6 stay finite.
+_TINY_EPSILON_RATIO = 1e-6
+
+SoftmaxResult = SolverResult
+
+
+@dataclass
+class SoftmaxSolver:
+    """Direct softmax weighting solver for subspace modes.
+
+    Computes transport weights via ``softmax(-C / epsilon)`` and scales
+    by the source marginal to produce a transport matrix whose row sums
+    equal ``a``. This is equivalent to entropic OT when the target marginal
+    constraint is naturally satisfied (few particles, subspace mode).
+
+    Attributes:
+        epsilon: Temperature parameter (entropic regularization strength).
+            Controls sharpness of the softmax: lower epsilon gives sharper
+            (more selective) weights.
+        compile: Placeholder for API compatibility with SinkhornSolver.
+            Currently unused since softmax is a single torch op.
+    """
+
+    epsilon: float = 0.1
+    compile: bool = False
+
+    def solve(
+        self,
+        cost_matrix: torch.Tensor,
+        a: Optional[torch.Tensor] = None,
+        b: Optional[torch.Tensor] = None,
+        init_f: Optional[torch.Tensor] = None,
+        init_g: Optional[torch.Tensor] = None,
+        scale_cost: Optional[Union[str, float]] = None,
+    ) -> SolverResult:
+        """Compute softmax weights from cost matrix.
+
+        Args:
+            cost_matrix: Cost matrix C of shape (P, V).
+            a: Source marginal of shape (P,). Defaults to uniform 1/P.
+            b: Target marginal (accepted but ignored for softmax).
+            init_f: Warm-start dual potential (accepted but ignored).
+            init_g: Warm-start dual potential (accepted but ignored).
+            scale_cost: Cost scaling: 'mean', 'max_cost', or float divisor.
+
+        Returns:
+            SolverResult with transport matrix, cost, and metadata.
+
+        Raises:
+            ValueError: If epsilon <= 0.
+        """
+        validate_positive(
+            self.epsilon,
+            "epsilon",
+            "epsilon is the temperature in softmax(-C/epsilon).",
+        )
+
+        P, V = cost_matrix.shape
+
+        # SoftmaxSolver is one-sided: it only enforces row sums equal `a`.
+        # If a caller passes a non-uniform `b` they probably meant to call
+        # SinkhornSolver; warn loudly so the constraint isn't silently dropped.
+        if b is not None:
+            uniform = torch.full_like(b, 1.0 / V)
+            if not torch.allclose(b.to(uniform.dtype), uniform, atol=1e-6):
+                warnings.warn(
+                    "SoftmaxSolver ignores the target marginal `b`: it only "
+                    "enforces row sums equal to the source marginal `a`. "
+                    "Pass solver='sinkhorn' for two-sided OT.",
+                    stacklevel=2,
+                )
+
+        # FP32 promotion (softmax subtracts the row max only after an outer
+        # autocast may have downcast -C/epsilon) + finite-cost handling
+        # (+Inf models a hard "never pick this vertex" constraint).
+        cost_matrix = sanitize_cost(cost_matrix)
+        device, dtype = cost_matrix.device, cost_matrix.dtype
+        a = align_marginal(a, P, device, dtype)
+
+        C = scale_cost_matrix(cost_matrix, scale_cost)
+        # softmax is shift-invariant, so recentering leaves the weights unchanged
+        # while min(C)=0 avoids a +inf logit that would NaN the softmax at tiny eps.
+        C, cost_shift = recenter_cost(C)
+
+        # ``epsilon > 0`` is enough to avoid division-by-zero, but
+        # eps=1e-30 with cost_max~10 still overflows -C/epsilon before
+        # softmax can subtract the row max. Warn at extreme ratios, but avoid a
+        # host sync every solve: only run the reduction when epsilon drops below
+        # the smallest value already checked (a new, sharper regime).
+        # Only runs on an epsilon decrease, so a fixed-epsilon cost blowup is not caught.
+        min_checked = getattr(self, "_min_eps_checked", None)
+        if min_checked is None or self.epsilon < min_checked:
+            self._min_eps_checked = self.epsilon
+            cost_max = C.detach().abs().max().item() if C.numel() > 0 else 0.0
+            if cost_max > 0 and self.epsilon < _TINY_EPSILON_RATIO * cost_max:
+                warnings.warn(
+                    f"SoftmaxSolver epsilon={self.epsilon:.2e} is very small "
+                    f"relative to the cost-matrix scale (max |C|={cost_max:.2e}); "
+                    f"-C/epsilon may underflow / overflow before the row-max "
+                    f"subtraction inside torch.softmax. Consider rescaling the "
+                    f"cost or raising epsilon.",
+                    stacklevel=2,
+                )
+
+        # PyTorch's softmax subtracts the row max internally for stability.
+        # Pin the whole block inside an autocast-disabled context so an outer
+        # mixed-precision region can't downcast intermediates back to BF16.
+        with torch.amp.autocast("cuda", enabled=False), torch.amp.autocast("cpu", enabled=False):
+            W = torch.softmax(-C / self.epsilon, dim=-1)
+
+            # Row sums equal source marginal a
+            transport = W * a.to(W.dtype).unsqueeze(-1)
+
+            # Undo the recenter shift so the reported cost is <C_scaled, transport>.
+            ent_cost = ((C * transport).sum() + cost_shift * a.to(W.dtype).sum()).item()
+
+        return SolverResult(
+            matrix=transport,
+            cost=ent_cost,
+            f=None,
+            g=None,
+            converged=True,
+            n_iters=1,
+            ent_reg_cost=ent_cost,
+        )