Transport 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 both transport surfaces.

Two surfaces

The torchmatch.transport namespace exposes two sub-packages that solve closely related problems with different input conventions and different runtime requirements:

• torchmatch.transport.matrix takes a precomputed cost matrix (shape (N, M) or (B, N, M)) and returns a transport plan in log space or a scalar divergence. Backed by Python custom ops that run iterative log-sum-exp loops on CPU and CUDA; EXACT_EMD uses a network-simplex algorithm (CPU only).
• torchmatch.transport.samples takes two raw point clouds (shapes (N, D) and (M, D)) and returns a scalar OT loss or debiased divergence. Pairwise squared distances are computed on the fly inside Triton kernels — no N×M matrix is materialised. CUDA only.

Both surfaces register their ops under torch.ops.transport.*; the Python-side dispatchers (transport.matrix.solve, transport.samples.loss) are the recommended entry points.

Input semantics

transport.matrix

• NaN cells are rejected with a RuntimeError.
• -inf cells are rejected with a RuntimeError. A cost of -inf would break the solver's numerical assumptions; use a large negative finite value if you need to strongly favour a particular pair.
• +inf cells are forbidden edges. The dispatcher passes them through to the LSE solver, where they cancel automatically; the network simplex handles them by withholding the edge from the graph.
• mask (optional bool tensor of the same shape as cost): False cells are rewritten to +inf before the solver runs.
• a and b (optional float tensors): source and target marginals. Both default to uniform; either can be supplied on its own. Both must be non-negative.

transport.matrix.solve raises ValueError for an out-of-range ndim, a non-float dtype, a non-CPU / non-CUDA device, a NaN or -inf cell, a non-positive reg, or an unknown backend string.

transport.samples

The samples surface is CUDA-only; loss(...) raises RuntimeError when called with CPU tensors. Triton kernels live under sources/torchmatch/transport/samples/kernels/ and load lazily on first call.

Tracing & graph capture

Both surfaces are compatible with torch.compile and support autograd. They can be used inside compiled model code or as differentiable loss functions.

See also

• Operations: the assignment ops, for comparison.
• Building: wheel vs JIT paths and environment overrides.