module

transport.samples.kernels.c_transform_sqeuclid

module torchmatch.transport.samples.kernels.c_transform_sqeuclid

C-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

Source: torchmatch/transport/samples/kernels/c_transform_sqeuclid.py:1