Optimal Transport

TorchMorph provides a differentiable Sinkhorn solver that runs on CPU or CUDA. CUDA float32 inputs automatically use the fused implementation.

torchmorph.build_cost_matrix

build_cost_matrix(shape, p=2, device=None)

Build pairwise distances between points on a spatial grid

The grid contains d = prod(shape) points in row-major order. Flatten grid-shaped distributions to (n, d) before passing this matrix to :class:SinkhornSolver.

Parameters:

Name	Type	Description	Default
shape	sequence[int]	Spatial grid shape.	required
p	float	Norm degree passed to :func:torch.cdist.	2
device	device or str	Device for the returned matrix.	None

Returns:

Type	Description
Tensor	torch.Tensor: Symmetric float32 cost matrix with shape (d, d).

Example

>>> import torchmorph as tm
>>> cost = tm.build_cost_matrix((2, 3))
>>> cost.shape
torch.Size([6, 6])
>>> cost[0, 1].item()
1.0

torchmorph.SinkhornSolver

SinkhornSolver(epsilon=1.0, max_iter=100, threshold=0.0, p=2, log_space=False)

Bases: Module

Solve entropy-regularized balanced optimal transport

The module accepts batches of flattened histograms and returns transport costs that are differentiable with respect to both marginals. CUDA float32 inputs use fused kernels; other device and dtype combinations use PyTorch operations. Set log_space=True for better stability at small regularization values.

Parameters:

Name	Type	Description	Default
epsilon	float	Positive entropy-regularization strength.	1.0
max_iter	int	Positive maximum number of Sinkhorn iterations.	100
threshold	float	Nonnegative early-stopping tolerance checked by the PyTorch implementation. Fused CUDA kernels always run max_iter.	0.0
p	float	Norm used by the default one-dimensional cost matrix.	2
log_space	bool	Run Sinkhorn iterations in the log domain.	False

Example

>>> import torch
>>> import torchmorph as tm
>>> source = torch.tensor([[1.0, 0.0, 0.0]])
>>> target = torch.tensor([[0.0, 0.0, 1.0]])
>>> solver = tm.SinkhornSolver(epsilon=1.0, max_iter=200)
>>> solver(source, target).shape
torch.Size([1])

data_preprocess

data_preprocess(source, target, cost_matrix=None)

Validate and normalize batched transport marginals

Negative values are clamped to zero and every row is normalized to unit mass. These operations remain in the autograd graph. If no cost matrix is supplied, the d bins are treated as points on a line.

Parameters:

Name	Type	Description	Default
source	Tensor	Floating-point source marginals with shape (n, d).	required
target	Tensor	Floating-point target marginals with the same shape as source.	required
cost_matrix	Tensor	Shared pairwise cost matrix with shape (d, d).	None

Returns:

Type	Description
Tensor	tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Normalized source,
Tensor	normalized target, and contiguous cost matrix on the source device
Tensor	and with the source dtype.

Example

>>> import torch
>>> import torchmorph as tm
>>> solver = tm.SinkhornSolver()
>>> a, b, cost = solver.data_preprocess(
...     torch.tensor([[1.0, 1.0]]), torch.tensor([[3.0, 1.0]])
... )
>>> a.sum().item(), b.sum().item(), cost.shape
(1.0, 1.0, torch.Size([2, 2]))

forward

forward(source, target, cost_matrix=None)

Compute transport costs between corresponding marginal rows

Parameters:

Name	Type	Description	Default
source	Tensor	Floating-point source marginals with shape (n, d).	required
target	Tensor	Floating-point target marginals with shape (n, d).	required
cost_matrix	Tensor	Shared cost matrix with shape (d, d). If None, uses a one-dimensional grid cost.	None

Returns:

Type	Description
Tensor	torch.Tensor: Differentiable transport costs with shape (n,).

Example

>>> import torch
>>> import torchmorph as tm
>>> solver = tm.SinkhornSolver(max_iter=200)
>>> solver(torch.ones(2, 4), torch.ones(2, 4)).shape
torch.Size([2])

plan

plan(source, target, cost_matrix=None)

Reconstruct transport plans in the log domain

This method runs without gradient tracking.

Parameters:

Name	Type	Description	Default
source	Tensor	Floating-point source marginals with shape (n, d).	required
target	Tensor	Floating-point target marginals with shape (n, d).	required
cost_matrix	Tensor	Shared cost matrix with shape (d, d).	None

Returns:

Type	Description
Tensor	torch.Tensor: Transport plans with shape (n, d, d).

Example

>>> import torch
>>> import torchmorph as tm
>>> solver = tm.SinkhornSolver(max_iter=200)
>>> solver.plan(torch.ones(1, 3), torch.ones(1, 3)).shape
torch.Size([1, 3, 3])

potentials

potentials(source, target, cost_matrix=None)

Compute centered dual potentials for both marginals

The potentials are the envelope-theorem gradients of the entropic transport cost with respect to normalized marginals. This method runs without gradient tracking.

Parameters:

Name	Type	Description	Default
source	Tensor	Floating-point source marginals with shape (n, d).	required
target	Tensor	Floating-point target marginals with shape (n, d).	required
cost_matrix	Tensor	Shared cost matrix with shape (d, d).	None

Returns:

Type	Description
Tensor	tuple[torch.Tensor, torch.Tensor]: Source and target potentials,
Tensor	each with shape (n, d) and zero mean along the last axis.

Example

>>> import torch
>>> import torchmorph as tm
>>> solver = tm.SinkhornSolver(max_iter=200)
>>> f, g = solver.potentials(torch.ones(1, 3), torch.ones(1, 3))
>>> f.shape, g.shape
(torch.Size([1, 3]), torch.Size([1, 3]))