transport.samples.kernels.c_transform_sqeuclid.c_transform_kernel
def c_transform_kernel(x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor, cost_scale: float = 1.0, allow_tf32: 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 = 3, autotune: bool = True) -> tuple[torch.Tensor, torch.Tensor]Compute streaming min + argmin over the factored cost.
Computes per source point i:
min_val_i = min_j [-coord_scale * dot(x_i, y_j) + bias_j]
min_idx_i = argmin_j [same]
This is the inner minimum only. The caller adds alpha_i = cost_scale * ||x_i||²
to get the full c-transform values.
Args:
x: Source coordinates [n, d], CUDA
y: Target coordinates [m, d], CUDA
bias: Pre-scaled bias [m] = cost_scale * ||y||² - ψ
cost_scale: Cost scaling (1.0 for ||x-y||², 0.5 for ||x-y||²/2)
allow_tf32: Enable TF32 for matmul
block_m, block_n, block_k: Manual block sizes (disables autotune)
num_warps: Number of warps (disables autotune)
num_stages: Pipeline stages
autotune: Enable autotuning
Returns:
min_vals: Inner minimum values [n], float32
argmin_idx: Argmin indices [n], int64 (cast from kernel int32)
Source: torchmatch/transport/samples/kernels/c_transform_sqeuclid.py:239