transport.samples.kernels.cg_dense.hvp_dense_cg
def hvp_dense_cg(x: torch.Tensor, y: torch.Tensor, f_hat: torch.Tensor, g_hat: torch.Tensor, A: torch.Tensor, eps: float, rho_x: float | None = None, rho_y: float | None = None, tau2: float = 1e-05, max_cg_iter: int = 100, cg_rtol: float = 1e-06, cg_atol: float = 1e-06) -> tuple[torch.Tensor, DenseCgInfo]Compute HVP using dense CG with cached transport plan.
This is a drop-in replacement for the streaming Triton HVP that is
8x faster for small problem sizes (n <= 8192).
Parameters
| Name | Type | Description |
|---|---|---|
| x | torch.Tensor | Source points (n, d) |
| y | torch.Tensor | Target points (m, d) |
| f_hat | torch.Tensor | Source raw-form potential (n,) |
| g_hat | torch.Tensor | Target raw-form potential (m,) |
| A | torch.Tensor | Input matrix for HVP (n, d) |
| eps | float | Entropy regularization |
| rho_x = None | float | None | Source marginal KL penalty (None = strict constraint) |
| rho_y = None | float | None | Target marginal KL penalty (None = strict constraint) |
| tau2 = 1e-05 | float | Tikhonov regularization (only used for balanced OT) |
| max_cg_iter = 100 | int | Maximum CG iterations |
| cg_rtol = 1e-06 | float | Relative CG tolerance |
| cg_atol = 1e-06 | float | Absolute CG tolerance |
Returns
torch.Tensor — Hessian-vector product H @ A, shape (n, d)
Source: torchmatch/transport/samples/kernels/cg_dense.py:207