method

transport.samples.kernels.cg_dense.CachedDenseHVP.__init__

def __init__(self, x: torch.Tensor, y: torch.Tensor, f_hat: torch.Tensor, g_hat: 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)

Initialize cached HVP context.

Parameters

NameTypeDescription
xtorch.TensorSource points (n, d)
ytorch.TensorTarget points (m, d)
f_hattorch.TensorSource raw-form potential (n,)
g_hattorch.TensorTarget raw-form potential (m,)
epsfloatEntropy regularization
rho_x = Nonefloat | NoneSource marginal KL penalty (None = balanced)
rho_y = Nonefloat | NoneTarget marginal KL penalty (None = balanced)
tau2 = 1e-05floatTikhonov regularization
max_cg_iter = 100intMaximum inner CG iterations
cg_rtol = 1e-06floatRelative CG tolerance
cg_atol = 1e-06floatAbsolute CG tolerance

Source: torchmatch/transport/samples/kernels/cg_dense.py:405