transport.samples.kernels.c_transform_sqeuclid
module torchmatch.transport.samples.kernels.c_transform_sqeuclidC-Transform (hard argmin) kernel for squared Euclidean cost.
Computes the non-entropic Kantorovich c-transform via streaming min + argmin:
c_i = min_j [cost_scale * ||x_i - y_j||² - ψ_j]
j*_i = argmin_j [cost_scale * ||x_i - y_j||² - ψ_j]
Factorization (same trick as the LSE kernel):
cost_scale * ||x-y||² - ψ = cost_scale*||x||² + (cost_scale*||y||² - ψ) - 2*cost_scale*(x·y)
= alpha_i + bias_j - coord_scale * dot(x_i, y_j)
where:
alpha_i = cost_scale * ||x_i||² (constant w.r.t. j, factors out of min)
bias_j = cost_scale * ||y_j||² - ψ_j
coord_scale = 2 * cost_scale
The kernel computes min_j[-coord_scale * dot(x_i, y_j) + bias_j] via tiled streaming.
The Python wrapper adds alpha_i back to get the final c-transform values.
Tie-breaking: across tiles, smallest-j wins (strict < comparison).
Within a tile, tl.argmin selects the first minimum per Triton lane ordering.
Kernel outputs int32 indices (saves SRAM/registers). Python wrapper casts to int64.
Members
function
- c_transform_kernelCompute streaming min + argmin over the factored cost.
Source: torchmatch/transport/samples/kernels/c_transform_sqeuclid.py:1