Getting started

torchmatch is a PyTorch extension providing two families of solvers:

• Assignment — the linear assignment problem (LAP): one-to-one matching that minimises total cost. Used in tracking-by-detection (assigning detector outputs to object tracks across frames), DETR-style set losses (matching model predictions to ground-truth targets before computing a training loss), and cluster evaluation (measuring how well predicted clusters align with ground-truth labels).
• Transport — optimal transport: computing the minimum-cost way to transform one probability distribution into another. Used for comparing point clouds, learning geometry-aware losses, and aligning feature distributions across datasets (domain adaptation).

Both problem families (assignment and transport) are registered as torch.ops.* custom ops (meaning they work natively with torch.compile and the rest of the PyTorch ecosystem).

Installation

pip install torchmatch

torchmatch ships prebuilt wheels for Python 3.13 with cu126, cu128, and cu130 CUDA variants. A CPU-only wheel is also available. If no prebuilt wheel matches your Python/CUDA version, pip falls back to the source distribution (sdist), which JIT-compiles the C++/CUDA extensions on first import (takes 30–90 s).

Requirements:

Importing

import torchmatch immediately loads both the assignment and transport sub-packages (no separate import needed for each):

import torch
import torchmatch

# Assignment: ready immediately
row_to_col = torchmatch.assignment.solve(torch.rand(8, 8))

# Transport matrix face: ready immediately
log_plan = torchmatch.transport.matrix.solve(torch.rand(8, 12))

# Transport samples face: CUDA required
x = torch.randn(512, 3, device='cuda')
y = torch.randn(512, 3, device='cuda')
loss = torchmatch.transport.samples.loss(x, y)

First assignment

Solve a single 8 × 8 cost matrix:

import torch
import torchmatch

cost = torch.tensor([
    [4.0, 1.0, 3.0],
    [2.0, 0.0, 5.0],
    [3.0, 2.0, 2.0],
])
row_to_col = torchmatch.assignment.solve(cost)      # [1, 0, 2]
print(cost[torch.arange(3), row_to_col].sum())      # 4.0, the optimum

Unmatched rows (rectangular input, N > M) return −1:

cost = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
print(torchmatch.assignment.solve(cost).tolist())   # [0, 1, -1]

See the Assignment quickstart for batched problems, unpacking, and direct op access.

First transport solve

Compute a transport plan between two distributions represented as a cost matrix (entropy-regularised, which makes the problem smooth and differentiable):

The result is returned in log space for numerical stability; call .exp() to recover the actual transport plan.

import torch
import torchmatch

cost = torch.rand(8, 12)                            # (source, target) cost
log_plan = torchmatch.transport.matrix.solve(cost)  # (8, 12) in log space
plan = log_plan.exp()
plan.backward(torch.ones_like(plan))               # differentiable

Compute an optimal-transport loss (Wasserstein distance) between two point clouds — a scalar measuring how far apart the two sets of points are as distributions (CUDA required):

x = torch.randn(512, 3, device='cuda', requires_grad=True)
y = torch.randn(512, 3, device='cuda')
loss = torchmatch.transport.samples.loss(x, y)
loss.backward()

See the Transport quickstart for regularisation strength, debiasing (a correction that removes entropy artifacts from the plan), and backend selection.

Op namespaces

Every op binds at two locations:

# Python-friendly (autocomplete, type-checks)
torchmatch.assignment.ops.jonker_dense(cost)
torchmatch.transport.matrix.ops.log_sinkhorn(cost, 0.1, 100, None, None, None, None)

# torch.ops namespace (inside torch.compile, dynamic dispatch)
torch.ops.assignment.jonker_dense(cost)
torch.ops.transport.log_sinkhorn(cost, 0.1, 100, None, None, None, None)

Both forms are identical objects: torchmatch.assignment.ops.jonker_dense is torch.ops.assignment.jonker_dense is True.

Development environment (Nix)

The repo ships a self-contained flake.nix:

nix develop                   # default = cu128
nix develop .#cpu             # CPU-only
nix develop .#cu126 / .#cu128 / .#cu130

Inside the shell: uv sync --extra cu128 --all-groups then nix run .#test.

Next steps

Assignment

• Quickstart — batched problems, unpacking, backends
• Tracking tutorial — end-to-end SORT-style tracking
• Algorithms — JV, Munkres, Lawler derivations
• Reference — op signatures and constraints
• Choosing — decision tree from benchmarks

Transport

• Quickstart — matrix.solve and samples.loss
• Point-cloud tutorial — Wasserstein training loss
• Algorithms — Sinkhorn, debiasing, unbalanced OT
• Reference — matrix.solve and samples.loss signatures
• Choosing — backend selection guide