transport.samples.kernels.apply_shifted
module torchmatch.transport.samples.kernels.apply_shiftedshifted-form apply kernels (shifted potentials, s_I cancellation).
This module contains streaming P @ V and P @ vec kernels that work with
SHIFTED potentials directly (f_hat = f - alpha, g_hat = g - beta), avoiding
the cost of converting between potential conventions.
Kernels:
- apply_plan_mat_shifted: P @ mat or P^T @ mat (2D grid, tiles over D)
- apply_plan_vec_shifted: P @ vec or P^T @ vec (1D grid, s_I cancellation)
Key insight: Compute score WITHOUT f_hat/g_hat in the tiled loop, then apply
row/column marginal correction at the end. This is the same numerical trick as
the online-softmax correction factor used in streaming attention kernels.
Members
function
- apply_plan_mat_shiftedApply P or P^T to a matrix using shifted potentials.
- apply_plan_vec_shiftedApply transport plan to vector using shifted-form with s_I cancellation.
Source: torchmatch/transport/samples/kernels/apply_shifted.py:1