transport.samples.kernels.streaming_sqeuclid.sinkhorn_lse
def sinkhorn_lse(x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor, eps: float, cost_scale: float = 1.0, damping: float = 1.0, allow_tf32: bool = True, use_exp2: bool = True, block_m: int | None = None, block_n: int | None = None, block_k: int | None = None, num_warps: int | None = None, num_stages: int = 2, autotune: bool = True) -> torch.TensorCompute shifted potential using the streaming-softmax kernel.
This computes: out_i = -ε * damping * LSE_j[coord_scale * x_i·y_j / ε + bias_j]
The kernel applies coord_scale = 2 * cost_scale inside the kernel to scale
the dot product, ensuring consistent TF32 rounding between fused and separate
kernel paths.
Args:
x: Source coordinates [n, d]
y: Target coordinates [m, d]
bias: Pre-scaled bias [m]
eps: Regularization parameter
cost_scale: Cost scaling (1.0 for full ||x-y||², 0.5 for half ||x-y||²/2)
damping: Unbalanced OT damping (1.0 for balanced)
allow_tf32: Enable TF32 for matmul
use_exp2: Use exp2/log2 for better numerical stability
block_m, block_n, block_k: Manual block sizes (disables autotune)
num_warps: Number of warps (disables autotune)
num_stages: Number of pipeline stages
autotune: Enable autotuning
Returns:
out: Shifted potential [n]
Source: torchmatch/transport/samples/kernels/streaming_sqeuclid.py:945