transport.samples.kernels.cg_python_batched.compiled_cg_solve
def compiled_cg_solve(x: torch.Tensor, y: torch.Tensor, f: torch.Tensor, g: torch.Tensor, diag_x: torch.Tensor, denom: torch.Tensor, rhs: torch.Tensor, eps: float, x2: torch.Tensor | None = None, y2: torch.Tensor | None = None, x0: torch.Tensor | None = None, max_iter: int = 50, rtol: float = 1e-06, atol: float = 1e-06, use_compile: bool = True) -> tuple[torch.Tensor, PythonBatchedCgInfo]Solve H @ x = b using CG with torch.compile optimization.
This version uses torch.compile for the core CG loop, providing
CUDA graph capture and kernel fusion for additional speedup.
Note: This runs a fixed number of iterations (no early exit) to be
torch.compile friendly. The convergence info reflects the state
after max_iter iterations.
Args:
x: Source points (n, d)
y: Target points (m, d)
f: Source potential (n,)
g: Target potential (m,)
diag_x: Diagonal D_x (n,) - row sums of P
denom: Denominator D_y (m,) - column sums of P plus regularization
rhs: Right-hand side b (m,)
eps: Regularization parameter
x2: Precomputed ||x||^2 (n,) [optional]
y2: Precomputed ||y||^2 (m,) [optional]
x0: Initial guess (m,) [optional, ignored in compiled version]
max_iter: Number of CG iterations to run (fixed, no early exit)
rtol, atol: Convergence tolerances (for info only, no early exit)
use_compile: Whether to use torch.compile (default True)
Returns:
sol: Solution x (m,)
info: Convergence information
Source: torchmatch/transport/samples/kernels/cg_python_batched.py:340