Operations reference

Signatures, parameter types, and per-op descriptions are auto-generated from the source and browseable in the API reference →.

This page covers cross-cutting rules that apply to every op.

Namespace

Every op binds under two equivalent names:

• torchmatch.assignment.ops.<op>: the Python attribute binding (autocomplete-friendly, importable via from torchmatch.assignment.ops import <op>).
• torch.ops.assignment.<op>: the PyTorch op-namespace form (use inside torch.compile regions or for dynamic dispatch by name).

Both names point to the same callable: torchmatch.assignment.ops.jonker_dense is torch.ops.assignment.jonker_dense returns True.

Outputs are int64 row→col tensors with -1 for unmatched rows.

Input semantics

These rules hold for every op under torch.ops.assignment.*, both through torchmatch.assignment.solve and through direct op calls:

• +inf is a forbidden edge (a row-column pair that must not be assigned). The op rewrites it internally to a large finite sentinel (max_finite + 1) * (K + 1) where K = max(rows, cols). The caller does not need to sanitise the tensor first.
• NaN is rejected with a RuntimeError. NaN signals an upstream bug (zero-norm cosine, singular Kalman covariance, log of zero); it is never a valid forbidden-edge marker.
• -inf is rejected with a RuntimeError. An unboundedly cheap edge would force the solver into an infinite loop, and -inf has no meaningful "must-not-assign" interpretation either.

torchmatch.assignment.solve raises ValueError for the same NaN / -inf conditions before reaching a backend; direct op calls raise the same class of error from the C++ entry point.

Output convention

For every op:

• Output dtype is int64 ("long").
• Single-problem ops return shape (N,). out[i] = j means row i matches column j. out[i] = -1 means row i is unmatched.
• Batched ops return shape (B, N) with the same per-row semantics.
• Matches into padded columns (from rectangular inputs in the JV ops) map to -1.

Tracing & graph capture

Every op registers a FakeTensor kernel (a shape-only stub used by torch.compile to trace the op without running real computation):

• The CUDA primed-zeros Hungarian ops — munkres, hybrid, and lawler — are cudagraph_unsafe. They perform host-side cudaStreamSynchronize calls to read managed-memory iteration flags; under torch.compile(mode="reduce-overhead") they trigger a graph break.
• jonker_dense_batch (CUDA backend) is fully capturable.
• The CPU ops carry no graph constraints; they do not participate in CUDA-graph capture.

Pure-Python functions

Two functions live outside torch.ops.assignment.* and are called directly from torchmatch.assignment:

auction_assignment

Bertsekas' synchronous auction algorithm — an iterative, pure-Python solver that works on CPU and CUDA without the compiled extension. Unlike the ops above, it returns a triple (matches, unmatched_rows, unmatched_cols) rather than a row→col tensor, and is not wired into solve.

torchmatch.assignment.auction_assignment(
    cost_matrix: Tensor,   # (N, M) float; +inf = forbidden; NaN/-inf rejected
    bid_size: float,       # bid step; epsilon = min(bid_size / min(N,M), 1e-3)
    max_iters: int = 100_000,
) -> tuple[Tensor, Tensor, Tensor]
# (K,2) matches, (N-K,) unmatched_rows, (M-K,) unmatched_cols

Integer and low-precision (float16, bfloat16) inputs are cast to float32 automatically. Raises RuntimeError if convergence is not reached within max_iters iterations.

assignment_cost

Computes the total cost of a LAP solution. Accepts the output of solve or the row→col column from auction_assignment's matches tensor.

torchmatch.assignment.assignment_cost(
    cost: Tensor,          # (N, M) or (B, N, M), float32/float64
    matches: Tensor,       # (N,) or (B, N), int64; -1 = unmatched
    *,
    reduction: str = "sum",  # "sum" | "mean" | "none"
) -> Tensor  # scalar / (B,) / (N,) / (B,N) depending on ndim and reduction

Unmatched rows (-1) contribute 0 to the result regardless of reduction.

See also

• Algorithms: why the two families exist.
• Benchmarks: per-op latency across cost distributions, with use-case recommendations.