transport.samples.kernels.streaming_sqeuclid.precompute_sinkhorn_inputs
def precompute_sinkhorn_inputs(x: torch.Tensor, y: torch.Tensor, a: torch.Tensor, b: torch.Tensor, eps: float, cost_scale: float = 1.0) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]Precompute static bias components.
Args:
x: Source points [n, d]
y: Target points [m, d]
a: Source marginal weights [n]
b: Target marginal weights [m]
eps: Regularization parameter
cost_scale: Cost scaling (1.0 for full ||x-y||², 0.5 for half ||x-y||²/2)
Returns:
alpha: Source squared norms [n] = cost_scale * ||x||²
beta: Target squared norms [m] = cost_scale * ||y||²
gamma: Scaled log source weights [n] = eps * log(a) (for g-update)
delta: Scaled log target weights [m] = eps * log(b) (for f-update)
Notes:
- For cost_scale=1.0: full squared Euclidean C = ||x-y||²
- For cost_scale=0.5: half squared Euclidean C = ||x-y||²/2
- Use sinkhorn_lse() with raw x, y coordinates (not pre-scaled Q, K)
Source: torchmatch/transport/samples/kernels/streaming_sqeuclid.py:69