Transport quickstart

import torchmatch also loads torchmatch.transport. Both sub-packages (transport.matrix and transport.samples) are ready as soon as the import returns; no separate install is required.

Your first matrix solve

torchmatch.transport.matrix.solve takes a cost matrix and returns a transport plan — a matrix P where Pi,j is the fraction of the total weight assigned from source row i to target column j, with entropic regularisation added to make the solution unique and differentiable. By default, the plan is computed using the log-domain Sinkhorn algorithm (see Algorithms for details).

import torch
import torchmatch

# Cost matrix: element (i, j) = cost of moving mass (probability weight, point density,
# or any quantity being redistributed) from source i to target j.
# Shape: (N, M) for a single problem; (B, N, M) for a batch.
cost = torch.rand(8, 12)

# Returns a log-plan of the same shape: (8, 12).
log_plan = torchmatch.transport.matrix.solve(cost)

# Exponentiate to get the transport plan in [0, 1].
plan = log_plan.exp()
print(plan.sum(dim=-1))   # ≈ uniform, each row sums to 1/N
print(plan.sum(dim=-2))   # ≈ uniform, each column sums to 1/M

The plan is fully differentiable: gradients flow back through log_plan.exp() and torchmatch.transport.matrix.solve to the cost matrix and, optionally, to the source and target weight vectors (marginals).

Regularisation

The reg parameter controls the entropic regularisation strength. Entropic regularisation adds a smoothness penalty that makes the plan unique and differentiable; larger reg spreads mass more evenly across all entries, while smaller reg concentrates it on the lowest-cost assignments (approaching the unregularised optimum).

# Sharp plan: closer to exact OT but slower to converge.
log_plan_sharp = torchmatch.transport.matrix.solve(cost, reg=0.01, n_iter=500)

# Soft plan: converges in few iterations; acts as a soft attention matrix.
log_plan_soft = torchmatch.transport.matrix.solve(cost, reg=1.0, n_iter=50)

Sinkhorn divergence (debiased scalar loss)

When you need a single scalar that measures how different two distributions are — rather than a full transport plan — use SINKHORN_DIVERGENCE. It removes a systematic bias present in the plain Sinkhorn loss, so the result is zero only when the two inputs are identical. It returns a scalar (or (B,) tensor for batches).

from torchmatch.transport.matrix import Backend

divergence = torchmatch.transport.matrix.solve(
    cost,
    backend=Backend.SINKHORN_DIVERGENCE,
    reg=0.1,
)
print(divergence)   # scalar ≥ 0; equals 0 when source == target
divergence.backward()

Your first point-cloud loss

torchmatch.transport.samples.loss takes two sets of points and returns a scalar that measures the optimal-transport distance between them. The pairwise distances are computed without materialising the full N×M cost matrix in memory, keeping peak memory low. CUDA only.

import torch
import torchmatch

x = torch.randn(512, 3, device='cuda', requires_grad=True)   # source: (N, D)
y = torch.randn(512, 3, device='cuda')                       # target: (M, D)

loss = torchmatch.transport.samples.loss(x, y)
print(loss)         # scalar ≥ 0

loss.backward()     # gradients flow through x (and y if requires_grad)
print(x.grad.shape) # (512, 3)

blur controls how spread out the matching is (analogous to the reg parameter in matrix.solve): lower values produce sharper, more concentrated matchings; higher values produce smoother ones.

# Lower blur → sharper matching; higher blur → smoother/faster
loss = torchmatch.transport.samples.loss(x, y, blur=0.01)

Pass debias=True to compute the Sinkhorn divergence (a symmetric, bias-corrected scalar distance) instead of the raw loss:

loss = torchmatch.transport.samples.loss(x, y, debias=True)

Calling an op directly

The individual ops are exported at torchmatch.transport.matrix.ops.* and at torch.ops.transport.*:

from torchmatch.transport.matrix.ops import log_sinkhorn, sinkhorn_divergence

a = torch.full((1, 8), 1.0 / 8)    # source marginal
b = torch.full((1, 12), 1.0 / 12)  # target marginal
cost_3d = cost.unsqueeze(0)         # (1, 8, 12)

log_plan = log_sinkhorn(cost_3d, 0.1, 200, a, b, None, None)

torchmatch.transport.matrix.ops.log_sinkhorn is torch.ops.transport.log_sinkhorn evaluates to True.

torch.compile

Both the matrix API and the samples API work under torch.compile:

@torch.compile
def compute_loss(x, y):
    return torchmatch.transport.samples.loss(x, y, debias=True)

The backward pass is registered so that torch.compile can fuse the forward and gradient computation into a single optimised graph.

Next steps

• Point-cloud tutorial: end-to-end shape-generation training.
• Algorithms: how Sinkhorn works, debiasing, and unbalanced OT.
• Reference: full API signatures for matrix.solve and samples.loss.
• Choosing: which backend to use for your problem.
• Assignment: if you need a hard, integer-valued, one-to-one matching rather than a soft plan, the assignment family is the right tool.