torchmatch 1.0.0

Linear assignment problem solvers for PyTorch.

torchmatch

Linear assignment problem (LAP) solvers for PyTorch, registered under torch.ops.matching. Inputs are torch.Tensor; outputs are integer row-to-column assignments.

📖 Documentation: https://khwstolle.github.io/torchmatch/. Tutorials, API reference, and the full benchmark report.

Two algorithmic families share the op surface:

• Munkres (Kuhn-Munkres 1957): the primed/starred-zeros 6-step augmenting-path procedure. CUDA only; three parallelization strategies.
• Jonker-Volgenant / successive-shortest-path (SSP): the Dijkstra-like family.

Both families return optimal assignments. They differ in the augmentation routine, the parallelization story, and constant factors.

Install

pip install torchmatch

The Munkres CUDA ops and the CUDA backend of jv_dense_batch need a CUDA-capable PyTorch build. CPU ops run on any platform; the SIMD-optimized variants need AVX2/FMA at build time (sdist falls back to the AVX2/FMA-aware JIT path, and jv_scalar works without SIMD).

Usage

import torch
import torchmatch                          # extensions load eagerly at import

cost = torch.randn(8, 8, device="cuda")
row_to_col = torchmatch.munkres_hybrid(cost)

cost_cpu = torch.randn(8, 8)
row_to_col_cpu = torchmatch.jv_scalar(cost_cpu)

The same ops live at torch.ops.matching.<op>, which is useful inside torch.compile regions or when dispatching by name.

For control over extension loading (preforked workers, latency-sensitive serving), the loaders stay exported and idempotent:

torchmatch.load_cpu()                      # already ran at import; no-op
torchmatch.load_cuda()

Every op returns an int64 row→col mapping of length N (the number of rows). A -1 entry marks an unmatched row.

Ops

Munkres family (CUDA only)

Primed/starred-zeros Hungarian/Munkres algorithm; three GPU parallelization strategies. All carry the cudagraph_unsafe tag (host-side syncs read managed-memory flags).

Op	dtype (internal)	Variant
munkres_classical	float32, column-major	Single-path augmentation per outer iteration; CUB segmented reductions for column-min
munkres_hybrid	float32, row-major	Adaptive: starts single-path, switches to tree augmentation when single-path stalls
munkres_tree	float64, row-major	Always-tree: parallel BFS finds all vertex-disjoint augmenting paths per outer iteration; cooperative-groups + Thrust scan

Jonker-Volgenant family, CPU (single problem)

Successive-shortest-path (Dijkstra-like over reduced costs). All accept rectangular (N, M) cost matrices in float32 or float64.

Op	Variant
jv_scalar	Sequential reference; no SIMD. Implements Crouse 2016. Rejects NaN / -inf; treats +inf as an infeasible edge.
jv_dense	AVX2 flat-pointer inner loop; rectangular-capable
jv_compact	AVX2-gather inner loop; square-only internal kernel (rectangular inputs padded by the wrapper)

Jonker-Volgenant family, batched

Op	Devices	Constraints
jv_dense_batch	CPU + CUDA	CPU: (B, N, M), any size, at::parallel_for over per-problem jv_dense. CUDA: shared-memory tiled kernel, one block per problem; requires (B, K, K) square with K ≤ MAX_TILE = 64.
jv_compact_batch	CPU	at::parallel_for over per-problem jv_compact
jv_dense_batch_unpacked	CPU	Returns (matches, unmatched_rows, unmatched_cols, n_matched), which saves a per-problem Python unpack
jv_compact_batch_unpacked	CPU	Same shape, compact variant

torch.compile / torch.export

Every op carries a FakeTensor kernel. The Munkres CUDA ops carry the cudagraph_unsafe tag (host-side syncs); the CUDA backend of jv_dense_batch is fully capturable.

Build modes

The per-device loaders prefer a prebuilt .so shipped in the wheel and fall back to JIT-compiling the C++/CUDA sources via torch.utils.cpp_extension.load. Both paths register the same torch.ops.matching.* ops; the choice is transparent to callers.

Building wheels

# default: builds CPU extension; builds CUDA extension if a CUDA
# toolchain is detected (torch.utils.cpp_extension.CUDA_HOME)
pip wheel . -w dist/

# CPU-only wheel
TORCHMATCH_SKIP_CUDA=1 pip wheel . -w dist/

# CUDA SM targets (default: PyTorch's current-device default)
TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0" pip wheel . -w dist/

The build system pairs setuptools with torch.utils.cpp_extension.BuildExtension. Sdists ship the full `` tree so the JIT path works on any CUDA-capable machine, even without a matching wheel.

Runtime overrides

• TORCHMATCH_FORCE_JIT=1: skip the prebuilt .so and recompile from source via cpp_extension.load. Useful for development and for diagnosing ABI mismatches.

Development environment

A self-contained flake.nix provides Python 3.13, uv, and a chosen CUDA toolkit (default 12.8) without touching your global system. With direnv:

direnv allow                      # picks devShells.default = cu128
NIX_DEVSHELL_NAME=cpu  direnv reload
NIX_DEVSHELL_NAME=cu130 direnv reload

Without direnv:

nix develop                       # default = cu128
nix develop .#cpu                 # CPU-only (sets TORCHMATCH_SKIP_CUDA=1)
nix develop .#cu126
nix develop .#cu130

Once inside the shell:

uv sync --all-groups              # populate .venv from pyproject.toml
uv run python -m pytest tests/    # run the test suite