transport.samples.kernels.streaming_sqeuclid.sinkhorn_symmetric_step
def sinkhorn_symmetric_step(x: torch.Tensor, y: torch.Tensor, f_hat: torch.Tensor, g_hat: torch.Tensor, log_a: torch.Tensor, log_b: torch.Tensor, eps: float, cost_scale: float = 1.0, alpha: float = 0.5, damping_f: float = 1.0, damping_g: float = 1.0, allow_tf32: bool = True, use_exp2: bool = True, autotune: 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, label_x: torch.Tensor | None = None, label_y: torch.Tensor | None = None, label_cost_matrix: torch.Tensor | None = None, lambda_x: float = 1.0, lambda_y: float = 0.0) -> tuple[torch.Tensor, torch.Tensor]Fused symmetric Sinkhorn step: computes both f and g updates in ONE kernel.
This is the key optimization over separate f/g kernel calls - reduces kernel
launch overhead by 50% and improves GPU occupancy.
KEY: Uses x, y directly (no Q, K pre-allocation). The coord_scale = 2*cost_scale
is applied inside the kernel, avoiding memory allocation overhead.
Args:
x: Source coordinates [n, d] (NOT pre-scaled!)
y: Target coordinates [m, d] (NOT pre-scaled!)
f_hat: Current shifted f potential [n]
g_hat: Current shifted g potential [m]
log_a: Log source weights [n]
log_b: Log target weights [m]
eps: Regularization parameter
cost_scale: Cost scaling (1.0 for full, 0.5 for half cost)
alpha: Averaging weight (0.5 for symmetric, 1.0 for full update)
damping_f: Unbalanced OT damping for f (1.0 for balanced)
damping_g: Unbalanced OT damping for g (1.0 for balanced)
allow_tf32: Enable TF32 for matmul
use_exp2: Use exp2/log2 optimization
autotune: Enable Triton autotuning (recommended)
block_m, block_n, block_k: Manual block sizes (disables autotune)
num_warps: Number of warps (disables autotune)
num_stages: Pipeline stages
label_x: int32/int64 labels for x [n] (OTDD)
label_y: int32/int64 labels for y [m] (OTDD)
label_cost_matrix: W [V, V] label distance matrix (OTDD)
lambda_x: Weight for Euclidean cost (default 1.0)
lambda_y: Weight for label cost (default 0.0 = no label cost)
Returns:
f_hat_new, g_hat_new: Updated shifted potentials
Source: torchmatch/transport/samples/kernels/streaming_sqeuclid.py:1251