assignment.sinkhorn_log_plan
def sinkhorn_log_plan(cost_matrix: torch.Tensor, epsilon: float = DEFAULT_EPSILON, num_iter: int = DEFAULT_NUM_ITER, row_marginal: torch.Tensor | None = None, col_marginal: torch.Tensor | None = None) -> torch.TensorCompute a log-domain Sinkhorn transport plan for a cost matrix.
Performs ``num_iter`` iterations of log-domain Sinkhorn updates on
the Gibbs kernel :math:`\log K = -C / \epsilon`, producing the log
of the entropy-regularized optimal transport plan. The result is
numerically stable with respect to ``inf`` entries in ``C`` (which
mark forbidden assignments) via ``torch.logsumexp``.
Parameters
| Name | Type | Description |
|---|---|---|
| cost_matrix | torch.Tensor | ``(N, M)`` cost matrix. ``inf`` entries are treated as forbidden assignments and produce ``-inf`` entries in the log-plan. |
| epsilon = DEFAULT_EPSILON | float | Entropy-regularization weight. |
| num_iter = DEFAULT_NUM_ITER | int | Number of Sinkhorn iterations. |
| row_marginal = None | torch.Tensor | None | ``(N,)`` target row marginal. Uniform ``1/N`` by default. |
| col_marginal = None | torch.Tensor | None | ``(M,)`` target column marginal. Uniform ``1/M`` by default. |
Returns
torch.Tensor — ``(N, M)`` tensor of ``log P`` where ``P`` is the Sinkhorn plan.
Raises
- ValueError — If ``epsilon`` is not strictly positive.
Source: unitrack/assignment/_soft.py:140