transport.samples.kernels.apply_shifted.apply_plan_vec_shifted
def apply_plan_vec_shifted(x: torch.Tensor, y: torch.Tensor, f_hat: torch.Tensor, g_hat: torch.Tensor, log_a: torch.Tensor, log_b: torch.Tensor, vec: torch.Tensor, eps: float, axis: int, cost_scale: float = 1.0, allow_tf32: bool = False, use_exp2: bool = True, block_m: int = 64, block_n: int = 64, block_k: int = 64, num_warps: int = 4, num_stages: int = 2) -> torch.TensorApply transport plan to vector using shifted-form with s_I cancellation.
Computes P @ vec (axis=1) or P^T @ vec (axis=0) where:
P[i,j] = a[i] * b[j] * exp((f_hat[i] + g_hat[j] + 2*cost_scale*x[i].y[j]) / eps)
Key optimization: The normalizing sum s_I cancels algebraically, so we use:
- axis=1: out_I = a_I * exp(f_hat_I/eps + m_I) * O_I
- axis=0: out_J = b_J * exp(g_hat_J/eps + m_J) * O_J
Args:
x: Source points [n, d], fp16/fp32
y: Target points [m, d], fp16/fp32
f_hat: Shifted source potential [n], fp32
g_hat: Shifted target potential [m], fp32
log_a: Log source weights [n], fp32
log_b: Log target weights [m], fp32
vec: Vector to apply [m] for axis=1, [n] for axis=0
eps: Regularization parameter
axis: 1 for P @ vec, 0 for P^T @ vec
cost_scale: Scaling for cost (0.5 for half_cost)
allow_tf32: Use TF32 for dot products
use_exp2: Use exp2 instead of exp
block_m, block_n, block_k: Block sizes
num_warps, num_stages: Triton tuning params
Returns:
out: Result vector [n] for axis=1, [m] for axis=0
Source: torchmatch/transport/samples/kernels/apply_shifted.py:961