Assignment quickstart

Setup

import torch
import torchmatch

import torchmatch eagerly loads the torchmatch.assignment sub-package (CPU extension, and the CUDA extension when a GPU is present), so torchmatch.assignment.solve and torchmatch.assignment.ops.* are ready as soon as the import returns.

A square problem

The same example, expressed two ways:

AUTO is the default backend selector: it inspects the problem size and device at call time and picks the fastest available op. The direct-op tab shows what you would call explicitly if you wanted to pin a specific kernel.

[1, 0, 2] reads as row 0 → col 1, row 1 → col 0, row 2 → col 2. Both entry points (solve and the direct op) return the same row→col mapping. They differ only in whether the caller chooses the kernel.

Three CPU variants

torchmatch.assignment.solve picks among jonker_scalar, jonker_compact, and jonker_dense for you. When you want to benchmark or pin a variant, call the op directly. All three return an optimal assignment; on tied cost matrices the specific optimum may differ between ops, but the total cost is identical.

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

for op_name in ("jonker_scalar", "jonker_dense", "jonker_compact"):
    op = getattr(torchmatch.assignment.ops, op_name)
    out = op(cost)
    total = cost[torch.arange(8), out].sum().item()
    print(f"{op_name:18s} -> {out.tolist()}  total={total:.4f}")

Pick by workload:

• jonker_scalar: sequential, no SIMD; consistent performance regardless of matrix size or cost distribution.
• jonker_dense: uses AVX2 SIMD via a flat memory layout; the default for any problem size or cost distribution.
• jonker_compact: uses AVX2 gather instructions; 20 to 30 % faster than dense on small-to-medium square problems whose costs vary gradually across rows (e.g. distance or IoU matrices, as opposed to sparse or integer-tied costs). See Choosing for the regime where each wins.

Rectangular cost matrices

Unmatched rows return -1.

Both entry points return a tensor of length nrows. When there are more rows than columns, the leftover rows cannot be matched and their entries are -1.

Forbidden edges with +inf

Mark cells as infeasible by setting them to +inf. The ops rewrite +inf to a per-call sentinel (a large finite stand-in value that any optimal solver avoids) internally; no preprocessing is needed.

When forbidden edges make it impossible to match every row to a distinct column, the solver reports those rows as -1.

NaN handling

NaN signals an upstream bug; both entry points reject it explicitly:

Both paths raise RuntimeError with a message explaining the NaN was found.

Empty inputs

Both entry points accept (0, 0), (N, 0), and (0, M) inputs:

Batched problems

Pass a 3-D tensor (B, N, M) to solve B independent problems in one call. solve returns (B, N).

import torch
import torchmatch

# 64 problems, each 32×32, on CPU
costs = torch.rand(64, 32, 32)
assignments = torchmatch.assignment.solve(costs)   # (64, 32)

# Check that each batch element got a valid assignment
assert (assignments >= 0).all()

The CPU dispatcher uses at::parallel_for to distribute problems across threads; the CUDA backend for jonker_dense_batch launches a tiled shared-memory kernel (square problems only, K ≤ 64).

Unpacked output

When you need matched pairs and unmatched sets separately, pass unpack=True:

matches, unmatched_rows, unmatched_cols, n_matched = torchmatch.assignment.solve(
    costs, unpack=True,
)
# matches[b, :n_matched[b]] = matched (row, col) pairs for batch b
# unmatched_rows[b], unmatched_cols[b] = unmatched indices

CUDA ops

The CUDA ops (munkres, lawler) use a different algorithm family from the CPU JV variants and are available when a GPU is present:

import torch
import torchmatch

cost = torch.rand(128, 128, device='cuda')
row_to_col = torchmatch.assignment.solve(cost)    # AUTO picks lawler at N=128

# Or pin a specific CUDA op
row_to_col = torchmatch.assignment.ops.munkres(cost)
row_to_col = torchmatch.assignment.ops.lawler(cost)

For batched CUDA, pass a 3-D square tensor with K ≤ 64:

costs_gpu = torch.rand(64, 32, 32, device='cuda')
assignments = torchmatch.assignment.solve(costs_gpu)   # uses jonker_dense_batch CUDA

The CUDA tiled jonker_dense_batch kernel is CUDA-graph-safe (compatible with CUDA graph capture for reduced kernel-launch overhead). munkres and lawler are not CUDA-graph-safe because they pause GPU execution to copy an iteration flag back to the CPU on every step.

What's next

• Batched tracking workflow: detection-to-track association on a batch of frames.
• Choosing the right op: the decision tree that solve automates.
• Transport: if you need to match probability distributions rather than discrete items, handle sources and destinations with different total weight, or produce a differentiable fractional plan instead of a hard one-to-one assignment, the transport family is the right tool.