torchmatch

A tracker reads detections each frame and assigns each one to a track. Build the per-frame (N_tracks × M_detections) cost matrix — typically 1 − IoU (intersection-over-union overlap) between predicted and detected bounding boxes, optionally filtered by distance between box centers (centroid gating) or visual similarity (appearance distance) — stack frames into a batch, and jonker_dense_batch solves the whole batch in one CUDA launch.

The _unpacked variant returns (matches, unmatched_tracks, unmatched_dets, n_matched) directly — removing the per-frame Python loop that most codebases use to separate matched from unmatched indices after getting the raw assignment.

DETR-style object detectors find the lowest-cost one-to-one pairing (the Hungarian assignment) between model predictions and ground-truth targets before computing the per-pair loss term. The matcher runs once per image every training step, so it dominates training wall time when called naively in a Python loop. Batch the per-image cost matrices into one (B, N_pred, N_gt) tensor; jonker_dense_batch solves them all at once.

When every problem is square (prediction count fixed, GT padded to the same count), jonker_compact_batch runs the tighter AVX2-gather inner loop on CPU. For mixed shapes, use jonker_dense_batch.

Clustering algorithms assign arbitrary label numbers, so "cluster 0" in one run may correspond to "cluster 2" in another. Metrics like Adjusted Rand Index therefore need to find the optimal one-to-one relabelling between predicted and ground-truth cluster IDs before counting agreements — that relabelling is exactly the assignment problem. Build the cost as the negative confusion-matrix entry; the linear assignment solver (LAP) returns the optimal cluster-to-cluster mapping.

Square dense costs at moderate K (≤ 256) are the regime jonker_compact was built for. For K ≥ 512, jonker_dense takes over.

When a linear assignment problem (LAP) appears in a tight Python loop and SciPy is the bottleneck, swap the call. jonker_dense produces the same optimal cost, accepts the same rectangular input, and stays inside torch so the cost matrix never round-trips through numpy. For a batch of independent LAPs, jonker_dense_batch runs at::parallel_for across problems on CPU or the tiled CUDA kernel on GPU.

For training generative models over point sets, transport.samples.loss(x, y) computes the Sinkhorn approximation of the Wasserstein distance between two point clouds — a measure of how much work it takes to move one distribution onto the other — without ever building the full N×M pairwise-cost matrix in memory. A Triton streaming kernel handles the computation directly in fast CUDA registers. Gradients flow analytically through both x and y; pass debias=True for the symmetric Sinkhorn divergence variant.

Hard one-to-one matching (like the Hungarian algorithm) is not differentiable — gradients cannot flow back through a discrete matching step. Replace it with a Sinkhorn plan: transport.matrix.solve(cost) returns a regularized transport plan T ∈ 0,1^{N×M} — fractional soft assignments stored in log-domain for numerical stability — that differentiate smoothly through the cost matrix. Reduce the regularization (reg) toward zero to sharpen the plan toward a hard permutation. Both 2-D (N, M) and batched 3-D (B, N, M) cost tensors are accepted.

Standard optimal transport requires the source and target distributions to have exactly equal total mass — every point must be fully accounted for. If one set contains outliers or the two domains differ in size, this constraint forces those outliers into the plan and corrupts the result. UNBALANCED_SINKHORN relaxes the marginal constraints via a KL-divergence penalty controlled by the reach (or reach_x, reach_y) parameter. Smaller reach is more lenient; reach → ∞ recovers balanced OT. Both the matrix and samples faces support unbalanced mode.