Backends and Batching

What AUTO does

torchmatch.assignment.solve(cost) defaults to backend=Backend.AUTO. At call time, it inspects the device, shape, and whether the matrix is square, then selects the fastest available op and dispatches. The resolution happens in Python at call time, not at graph-capture time, so the chosen op traces cleanly under torch.compile.

You do not need to think about backends for everyday use. This tutorial covers them so you understand what solve is selecting, and to show when pinning a specific backend is worthwhile.

The three CPU backends

All three CPU backends implement the Jonker-Volgenant (JV) successive shortest-path algorithm. They differ in their inner loop:

• jonker_scalar is the sequential reference implementation, with no SIMD. It is the most portable: it runs correctly on any CPU, including those without AVX2. AUTO selects it when the problem is tiny (N * M <= 64), where the overhead of setting up SIMD outweighs the benefit.
• jonker_dense uses AVX2 vector instructions over a flat contiguous memory layout. It handles rectangular matrices natively (the underlying algorithm from Crouse 2016 supports n != m without padding). AUTO picks this as the default for any CPU problem above the scalar threshold.
• jonker_compact also uses AVX2 but with a tighter address pattern suited to square problems. On square matrices at N <= 256 with smoothly distributed costs (uniform random, gamma-distributed scores, L2 distances), it runs 20 to 30 percent faster than jonker_dense. On IoU costs or gated-sparse matrices it falls behind. It is square-only; calling it on a rectangular matrix pads internally.

import torch
import torchmatch

torch.manual_seed(0)
cost = torch.rand(64, 64, dtype=torch.float64)

for name in ("jonker_scalar", "jonker_dense", "jonker_compact"):
    op = getattr(torchmatch.assignment.ops, name)
    out = op(cost)
    total = cost[torch.arange(64), out].sum().item()
    print(f"{name:20s} total={total:.4f}")
# All three produce the same optimal total; tied costs may yield different
# but equally optimal assignments.

The CUDA backends

Two CUDA backends implement a different algorithm family: the Hungarian primal-dual method with starred and primed zeros.

• munkres follows Munkres' 1957 single-path procedure. It is fastest for small problems (N < 32) and for integer-tied cost matrices, where many cells share the same value and augmenting paths are short. It is also the only CUDA op that beats the CPU JV variants, specifically on integer-tied costs at large N.
• lawler uses Lawler's 1976 tree-augmentation variant, which finds all vertex-disjoint augmenting paths in a single breadth-first pass instead of one at a time. This exposes more GPU parallelism and makes lawler faster than munkres for large, fully-populated cost matrices (N >= 512).

Both require the cost tensor on a CUDA device. Neither supports CUDA graph capture, because they synchronize with the CPU between iterations; for graph-safe batched solving, see the next section.

Quick decision table

Pinning a backend directly

To skip AUTO and call a specific op, use torchmatch.assignment.ops:

cost = torch.rand(128, 128)
row_to_col = torchmatch.assignment.ops.jonker_dense(cost)

cost_gpu = cost.to("cuda")
row_to_col_gpu = torchmatch.assignment.ops.lawler(cost_gpu)

Pinning is useful for benchmarking and for making the backend choice visible in the source.

Batched solving with a 3-D tensor

Many workloads solve many independent assignment problems in parallel (one per frame in a video, one per anchor in a detector head). Pass a 3-D tensor of shape (B, N, M):

torch.manual_seed(42)
costs = torch.rand(32, 48, 48)       # 32 independent 48x48 problems
assignments = torchmatch.assignment.solve(costs)
print(assignments.shape)             # torch.Size([32, 48])
print((assignments >= 0).all())      # True: all rows matched (square)

On CPU, the dispatcher uses PyTorch's parallel_for to distribute problems across threads. On CUDA, a tiled shared-memory kernel solves one problem per CUDA block; this is restricted to square inputs with K <= 64.

costs_gpu = torch.rand(64, 32, 32, device="cuda")
assignments_gpu = torchmatch.assignment.solve(costs_gpu)  # uses jonker_dense_batch CUDA
print(assignments_gpu.shape)   # torch.Size([64, 32])

For K > 64 on CUDA, move the tensor to CPU for the solve and transfer the result back:

costs_large = torch.rand(16, 128, 128, device="cuda")
assignments = torchmatch.assignment.solve(costs_large.cpu()).to("cuda")

Unpacked output: matched pairs and unmatched sets

The default batch output is shape (B, N): entry [b, i] is the column assigned to row i in problem b, or -1 if unmatched. Pass unpack=True to get matched pairs and unmatched sets in one call without a Python loop:

costs = torch.rand(16, 40, 40)
matches, unmatched_rows, unmatched_cols, n_matched = torchmatch.assignment.solve(
    costs, unpack=True,
)
# matches[b, k]         : (row, col) pair of the k-th match in problem b
# unmatched_rows[b, k]  : index of the k-th unmatched row in problem b
# unmatched_cols[b, k]  : index of the k-th unmatched column in problem b
# n_matched[b]          : actual number of matches in problem b (rest are padding -1)

# Use a specific batch element:
b = 0
nm = int(n_matched[b].item())
real_matches = matches[b, :nm]           # shape (nm, 2)
lost_rows    = unmatched_rows[b, :costs.size(1) - nm]
new_cols     = unmatched_cols[b, :costs.size(2) - nm]

The _unpacked variants cost about 5 percent more than the packed variants. Use them whenever you would otherwise iterate over the batch in Python to compute the unmatched sets.

GPU vs CPU: the practical guideline

For small to medium problems (N <= 256), CPU is almost always faster than CUDA for assignment solving. The JV inner loop runs sequentially per problem, and a modern AVX2 core finishes it before the GPU has moved the first iteration flag to host memory. CUDA wins on integer-tied costs (where munkres often beats even jonker_dense on CPU) and for batched square problems with K <= 64 using the tiled kernel (the only CUDA-graph-safe option). The Choosing page has the full benchmark-backed decision tree.

See also

• Choosing the right op: the decision tree behind AUTO, with benchmark numbers.
• Reference: exact input/output shapes and constraints for every op.