module

transport.samples.kernels.apply_shifted

module torchmatch.transport.samples.kernels.apply_shifted

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

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