transport.samples.kernels.apply_raw.apply_plan_mat_sqeuclid
def apply_plan_mat_sqeuclid(x: torch.Tensor, y: torch.Tensor, f: torch.Tensor, g: torch.Tensor, mat: torch.Tensor, eps: float, axis: int, cost_scale: float = 1.0, scale: torch.Tensor | None = None, 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, 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 Pt to a matrix without materializing P (streaming, stable).
.. deprecated::
Use :func:`apply_plan_mat_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) * mat[j, :]
axis=0: out[j, :] = sum_i exp((f_i + g_j - cost_scale*C_ij)/eps) * mat[i, :]
If ``scale`` is provided (axis=1 only), the kernel uses `mat[j,:] * scale[j]`
on the fly (avoids allocating `mat * scale[:,None]`).
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:342