assignment.SoftAssignment.solve_with_plan
def solve_with_plan(self, cost_matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]Solve and return matches together with the Sinkhorn log-plan.
Threads the Sinkhorn log-plan back to the caller in one solve.
Use this when a downstream consumer (e.g. :class:`~unitrack.states.SoftReplace`)
needs the same plan that produced the matches — calling
:meth:`~unitrack.assignment.Assignment.forward` and re-running
:func:`sinkhorn_log_plan` would
compute the plan twice on different inputs (the threshold-masked
cost vs the un-masked original).
Parameters
| Name | Type | Description |
|---|---|---|
| cost_matrix | torch.Tensor | ``(N, M)`` cost matrix. |
Returns
torch.Tensor — ``(K, 2)`` long tensor of matched ``(row, col)`` indices.
Source: unitrack/assignment/_soft.py:82