transport.samples.kernels.streaming_sqeuclid.sinkhorn_lse_fused
def sinkhorn_lse_fused(x: torch.Tensor, y: torch.Tensor, g_hat: torch.Tensor, log_w: 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.TensorFused LSE kernel that computes bias in SRAM (matches symmetric kernel interface).
This computes: out_i = -ε * damping * LSE_j[coord_scale * x_i·y_j / ε + ĝ_j/ε + log(w_j)]
KEY OPTIMIZATION: Load ĝ and log(w) separately and compute bias = ĝ/ε + log(w) in SRAM.
This matches the symmetric kernel interface and eliminates Python kernel launch overhead.
Args:
x: Source coordinates [n, d]
y: Target coordinates [m, d]
g_hat: Shifted potential [m] (ĝ = g - β for f-update)
log_w: Log marginal [m] (log(b) for f-update, NOT scaled by eps!)
eps: Regularization parameter
cost_scale: Cost scaling (1.0 for full, 0.5 for half)
damping: Unbalanced OT damping (1.0 for balanced)
allow_tf32: Enable TF32 for matmul
use_exp2: Use exp2/log2 for numerical stability
block_m, block_n, block_k, num_warps: Manual block sizes
num_stages: Pipeline stages
autotune: Enable autotuning
Returns:
out: Shifted potential [n]
Source: torchmatch/transport/samples/kernels/streaming_sqeuclid.py:1101