src.lap_solvers.hungarian.hungarian

src.lap_solvers.hungarian.hungarian(s: torch.Tensor, n1: Optional[torch.Tensor] = None, n2: Optional[torch.Tensor] = None, nproc: int = 1) → torch.Tensor[source]

Solve optimal LAP permutation by hungarian algorithm. The time cost is \(O(n^3)\).

Parameters

  • s – \((b\times n_1 \times n_2)\) input 3d tensor. \(b\): batch size

  • n1 – \((b)\) number of objects in dim1

  • n2 – \((b)\) number of objects in dim2

  • nproc – number of parallel processes (default: nproc=1 for no parallel)

Returns

\((b\times n_1 \times n_2)\) optimal permutation matrix

Note

We support batched instances with different number of nodes, therefore n1 and n2 are required to specify the exact number of objects of each dimension in the batch. If not specified, we assume the batched matrices are not padded.