assignment.soft_assignment
def soft_assignment(cost_matrix: torch.Tensor, epsilon: float = DEFAULT_EPSILON, num_iter: int = DEFAULT_NUM_ITER) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]Derive a discrete assignment from a Sinkhorn transport plan.
Runs :func:`.sinkhorn_log_plan` on the cost matrix, then extracts a
full discrete assignment by running the Hungarian algorithm on
``-log_plan`` — equivalently, maximising the plan's log-likelihood
over a bipartite matching. Pairs whose underlying cost is non-finite
are rejected and moved to the residual.
The hard extraction is non-differentiable. When training with a
soft-assignment loss, call :func:`.sinkhorn_log_plan` directly and
compute the loss from the returned log-plan.
Parameters
| Name | Type | Description |
|---|---|---|
| cost_matrix | torch.Tensor | ``(N, M)`` cost matrix. ``inf`` entries mark forbidden pairs. |
| epsilon = DEFAULT_EPSILON | float | Entropy-regularization weight, forwarded to Sinkhorn. |
| num_iter = DEFAULT_NUM_ITER | int | Number of Sinkhorn iterations, forwarded to Sinkhorn. |
Returns
torch.Tensor — ``(K, 2)`` long tensor of matched ``(row, col)`` indices.
Source: unitrack/assignment/_soft.py:202