transport.samples.kernels.apply_shifted.apply_plan_mat_shifted
def apply_plan_mat_shifted(x: torch.Tensor, y: torch.Tensor, f_hat: torch.Tensor, g_hat: torch.Tensor, log_a: torch.Tensor, log_b: torch.Tensor, mat: torch.Tensor, eps: float, axis: int, cost_scale: float = 1.0, scale: torch.Tensor | None = None, block_m: int | None = None, block_n: int | None = None, block_k: int | None = None, block_d: int | None = None, num_warps: int = 4, num_stages: int = 2, use_exp2: bool = True, allow_tf32: bool = False, autotune: bool = True) -> torch.TensorApply P or P^T to a matrix using shifted potentials.
This kernel works with SHIFTED potentials directly:
f_hat = f - alpha where alpha = cost_scale * ||x||^2
g_hat = g - beta where beta = cost_scale * ||y||^2
Computes:
axis=1: out[i, :] = sum_j P_ij * mat[j, :] (P @ V)
axis=0: out[j, :] = sum_i P_ij * mat[i, :] (P^T @ V)
where P_ij = a_i * b_j * exp((f_hat + g_hat + 2*cost_scale*x.y) / eps)
Args:
x: Source points [n, d]
y: Target points [m, d]
f_hat: Shifted f potential [n] (f_hat = f - cost_scale * ||x||^2)
g_hat: Shifted g potential [m] (g_hat = g - cost_scale * ||y||^2)
log_a: Log source weights [n]
log_b: Log target weights [m]
mat: Matrix V [m, d] for axis=1, [n, d] for axis=0
eps: Regularization parameter
axis: 1 for P @ V, 0 for P^T @ V
cost_scale: Scaling for cost (1.0 for ||x-y||^2, 0.5 for ||x-y||^2/2)
scale: Optional per-row scale [m] for axis=1 (applies mat * scale[:, None])
autotune: If True (default), use autotuned kernel configs
Returns:
out: Result matrix [n, d] for axis=1, [m, d] for axis=0
Note:
The shifted potentials can be obtained from standard potentials via:
f_hat = f - cost_scale * (x ** 2).sum(dim=1)
g_hat = g - cost_scale * (y ** 2).sum(dim=1)
Or directly from the shifted-potential Sinkhorn solvers.
Source: torchmatch/transport/samples/kernels/apply_shifted.py:387