transport.samples.loss
def loss(x: torch.Tensor, y: torch.Tensor, blur: float = 0.05, debias: bool = False, reach: float | None = None, reach_x: float | None = None, reach_y: float | None = None, a: torch.Tensor | None = None, b: torch.Tensor | None = None, scaling: float = 0.5, n_iter: int | None = None, threshold: float | None = None, half_cost: bool = False, p: int = 2) -> torch.TensorScalar OT cost / divergence over point clouds (CUDA only).
Parameters
| Name | Type | Description |
|---|---|---|
| x | torch.Tensor | Source point cloud (n, d), float32, CUDA. |
| y | torch.Tensor | Target point cloud (m, d), float32, CUDA. |
| blur = 0.05 | float | Bandwidth parameter; the entropy regularization is ``eps = blur**2``. |
| debias = False | bool | If True, returns the Sinkhorn divergence ``S_eps(x,y) - 0.5*S_eps(x,x) - 0.5*S_eps(y,y)`` which vanishes when ``x == y`` (unbiased). Requires three Sinkhorn solves. |
| reach = None | float | None | Unbalanced OT: KL marginal penalty ``rho = reach**2`` applied to both source and target. ``None`` = balanced OT. |
| reach_x = None | float | None | Semi-unbalanced: KL penalty for the source marginal only. |
| reach_y = None | float | None | Semi-unbalanced: KL penalty for the target marginal only. |
| a = None | torch.Tensor | None | Source weights (n,). Uniform if None. |
| b = None | torch.Tensor | None | Target weights (m,). Uniform if None. |
| scaling = 0.5 | float | Geometric decay factor for the epsilon schedule, in ``(0, 1)``. Smaller values converge faster but may be less numerically stable. |
| n_iter = None | int | None | Not supported; always raises. Control iterations via ``scaling``. |
| threshold = None | float | None | Early-stopping threshold on potential change. Incompatible with ``torch.compile`` (forces a host sync each check). |
| half_cost = False | bool | If True, uses ``cost = 0.5 * ||x - y||²`` instead of ``||x - y||²``. |
| p = 2 | int | Cost exponent. Only ``p=2`` is supported. |
Returns
loss — Scalar OT cost (shape ``()`` for 2-D input; ``(B,)`` for 3-D batch).
Source: torchmatch/transport/samples/_loss.py:52