transport.samples.kernels.apply_raw.mat5_sqeuclid
def mat5_sqeuclid(x: torch.Tensor, y: torch.Tensor, f: torch.Tensor, g: torch.Tensor, A: torch.Tensor, eps: float, cost_scale: float = 1.0, x2: torch.Tensor | None = None, y2: 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.TensorCompute Mat5 term for HVP: Mat5 = (-4*cost_scale/eps) * sum_j P_ij (A_i.y_j) y_j.
Args:
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:492