transport.samples.kernels.apply_raw.apply_plan_vec_sqeuclid
def apply_plan_vec_sqeuclid(x: torch.Tensor, y: torch.Tensor, f: torch.Tensor, g: torch.Tensor, vec: torch.Tensor, eps: float, axis: int, cost_scale: float = 1.0, x2: torch.Tensor | None = None, y2: torch.Tensor | None = None, log_a: torch.Tensor | None = None, log_b: torch.Tensor | None = None, block_m: int | None = None, block_n: int | None = None, block_k: 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 Pt to a vector without materializing P (streaming, stable).
.. deprecated::
Use :func:`apply_plan_vec_shifted` instead for better performance
with shifted-form shifted potentials.
Computes:
axis=1: out[i] = sum_j exp((f_i + g_j - cost_scale*C_ij)/eps) * vec[j]
axis=0: out[j] = sum_i exp((f_i + g_j - cost_scale*C_ij)/eps) * vec[i]
Args:
f: raw-form source potential [n] (includes absorbed log marginal)
g: raw-form target potential [m] (includes absorbed log marginal)
log_a: Optional log source weights [n]. If None, assumes uniform (log(1/n)).
log_b: Optional log target weights [m]. If None, assumes uniform (log(1/m)).
cost_scale: Scaling for cost function. 1.0 for full ||x-y||^2, 0.5 for half.
autotune: If True (default), use autotuned kernel configs for best performance.
If False, use manual block sizes (useful for reproducible benchmarks).
Source: torchmatch/transport/samples/kernels/apply_raw.py:221