API Reference

Solver

torchgw.sampled_gw(X_source=None, X_target=None, p=None, q=None, *, distance_mode='dijkstra', dist_source=None, dist_target=None, k=30, n_landmarks=50, fgw_alpha=0.0, C_linear=None, M=50, alpha=0.9, epsilon=0.001, s_shared=None, lambda_ema_beta=None, max_iter=500, tol=1e-05, min_iter_before_converge=50, mixed_precision=False, T_init=None, device=None, verbose=False, verbose_every=20, log=False, semi_relaxed=False, rho_a=1.0, rho_b=1.0, rho=None, multiscale=False, n_coarse=None, differentiable=False, grad_mode='implicit')[source]

Sampled Gromov-Wasserstein alignment between two datasets.

Standard log-domain Sinkhorn with slack augmentation for partial transport.

Parameters:

  • X_source (ndarray or Tensor, optional) – Feature matrices (N, D) and (K, D’). Required unless distance matrices or C_linear are given.

  • X_target (ndarray or Tensor, optional) – Feature matrices (N, D) and (K, D’). Required unless distance matrices or C_linear are given.

  • p (ndarray or Tensor, optional) – Marginal distributions (uniform if None).

  • q (ndarray or Tensor, optional) – Marginal distributions (uniform if None).

  • **Distances**

  • distance_mode (str) – "dijkstra" (default), "precomputed", or "landmark".

  • dist_source (ndarray or Tensor, optional) – Precomputed (N, N) and (K, K) distance matrices.

  • dist_target (ndarray or Tensor, optional) – Precomputed (N, N) and (K, K) distance matrices.

  • k (int) – k-NN graph connectivity.

  • n_landmarks (int) – Landmark count for distance_mode="landmark".

  • GW** (**Fused)

  • fgw_alpha (float) – 0 = pure GW, 1 = pure Wasserstein, in between = fused.

  • C_linear (ndarray or Tensor, optional) – (N, K) feature cost matrix. Required when fgw_alpha > 0.

  • algorithm** (**Sampled GW)

  • M (int) – Anchor pairs sampled per iteration.

  • alpha (float) – Momentum: T ← (1-alpha) T + alpha T_new.

  • epsilon (float) – Entropic regularization (decays 10x over the run).

  • s_shared (int, optional) – Expected number of matched points (partial transport).

  • lambda_ema_beta (float, optional) – EMA smoothing of the sampled cost matrix (Lambda ← (1-beta) Lambda + beta Lambda_sample). Reduces sampling variance at the cost of a small bias that vanishes at convergence. Typical values 0.3–0.7. None disables (default).

  • control** (**Solver)

  • max_iter (int, float) – Iteration cap and plan-change tolerance.

  • tol (int, float) – Iteration cap and plan-change tolerance.

  • min_iter_before_converge (int) – Iterations before convergence checks start.

  • mixed_precision (bool) – Run Sinkhorn in float32 (log domain, numerically safe), output float64.

  • T_init (ndarray or Tensor, optional) – Initial transport plan (overrides multiscale).

  • device (torch.device, optional)

  • verbose (bool, int)

  • verbose_every (bool, int)

  • log (bool) – Also return a dict with err_list, n_iter, gw_cost.

  • transport** (**Unbalanced)

  • semi_relaxed (bool) – Relax marginals via KL penalties.

  • rho_a (float) – KL penalty weights for the source and target marginals; smaller relaxes more. rho is the legacy alias for rho_b.

  • rho_b (float) – KL penalty weights for the source and target marginals; smaller relaxes more. rho is the legacy alias for rho_b.

  • start** (**Multiscale warm)

  • multiscale (bool) – Solve a coarse problem first and upsample its plan as T_init.

  • n_coarse (int, optional) – Coarse problem size (auto if None).

  • mode** (**Differentiable)

  • differentiable (bool) – Keep the computation graph so gradients flow to C_linear. Requires fgw_alpha > 0; use epsilon >= 0.01.

  • grad_mode (str) – "implicit" (default): exact gradient via the adjoint system at the Sinkhorn fixed point, memory O(NK + K^2). "unrolled": backprop through the Sinkhorn iterations, memory O(NK · iterations).

  • rho (float | None)

Returns:

  • T (Tensor (N, K))

  • log_dict (dict (only if log=True))

Return type:

Tensor | tuple[Tensor, dict]

torchgw.sampled_lowrank_gw(X_source=None, X_target=None, p=None, q=None, *, rank=20, lr_max_iter=5, lr_dykstra_max_iter=50, distance_mode='dijkstra', dist_source=None, dist_target=None, k=30, n_landmarks=50, fgw_alpha=0.0, C_linear=None, M=50, alpha=0.9, epsilon=0.001, s_shared=None, lambda_ema_beta=None, max_iter=500, tol=1e-05, min_iter_before_converge=50, mixed_precision=False, device=None, verbose=False, verbose_every=20, log=False, semi_relaxed=False, rho_a=1.0, rho_b=1.0, rho=None, multiscale=False, n_coarse=None)[source]

Sampled Gromov-Wasserstein with low-rank Sinkhorn.

Uses the low-rank factorization of Scetbon, Cuturi & Peyré (2021) to cut memory from O(NK) to O((N+K)·rank) per Sinkhorn step. This is a memory optimization for N, K > 50k; at smaller scales sampled_gw is faster.

Parameters:
Return type:

Tensor | tuple[Tensor, dict]

:param All other parameters are as in sampled_gw(). semi_relaxed and: :param differentiable are not supported by the low-rank solver.:

Returns:

  • T (Tensor (N, K))

  • log_dict (dict (only if log=True))

Parameters:
Return type:

Tensor | tuple[Tensor, dict]

Graph Construction

torchgw.build_knn_graph(X, k=30)[source]

Build a k-NN graph, stitching disconnected components if needed.

Parameters:

  • X (ndarray of shape (n_samples, n_features))

  • k (number of nearest neighbors)

Returns:

graph – Symmetric sparse distance graph, guaranteed connected.

Return type:

csr_matrix of shape (n_samples, n_samples)

Joint Embedding

torchgw.joint_embedding(anchor_name, data_by_name, graphs_by_name, transport_plans, lambda_reg=1.0, out_dim=30)[source]

Compute joint manifold embedding using transport plans.

Parameters:

  • anchor_name (str) – Name of the anchor/reference dataset.

  • data_by_name (dict mapping name -> (n_samples, n_features) array)

  • graphs_by_name (dict mapping name -> kNN graph (csr_matrix))

  • transport_plans (dict mapping (query_name, anchor_name) -> T array)

  • lambda_reg (float) – Regularization weight for Laplacian.

  • out_dim (int) – Dimensionality of output embedding.

Returns:

embeddings

Return type:

dict mapping name -> (n_samples, out_dim) array

Internal Modules

Sinkhorn

torchgw._sinkhorn.sinkhorn(a, b, C, reg, *, max_iter=100, tol=0.0005, check_every=10, tau_a=1.0, tau_b=1.0, inplace_C=False, verbose=False, log_u_init=None, log_v_init=None)[source]

Log-domain Sinkhorn without autograd.

Parameters:

  • a ((N,), (K,) marginals)

  • b ((N,), (K,) marginals)

  • C ((N, K) cost matrix. Overwritten with -C / reg if inplace_C) – (saves one N x K buffer; the caller must not reuse C).

  • reg (entropic regularization)

  • tau_a (KL damping per side, see kl_damping())

  • tau_b (KL damping per side, see kl_damping())

  • log_u_init (warm-start potentials from a previous call)

  • log_v_init (warm-start potentials from a previous call)

  • max_iter (int)

  • tol (float)

  • check_every (int)

  • inplace_C (bool)

  • verbose (bool)

Returns:

  • T ((N, K) transport plan)

  • log_u, log_v (dual potentials (detached))

Return type:

tuple[Tensor, Tensor, Tensor]

torchgw._sinkhorn.sinkhorn_differentiable(a, b, C, reg, *, max_iter=100, tol=0.0005, check_every=10, tau_a=1.0, tau_b=1.0, grad_mode='implicit', log_u_init=None, log_v_init=None)[source]

Sinkhorn with autograd support. Same contract as sinkhorn().

grad_mode

"implicit": exact VJP via the adjoint system at the fixed point, memory O(NK + K^2). "unrolled": backprop through the iterations, memory O(NK · iterations). "approximate": frozen potentials, inexact.

Parameters:
Return type:

tuple[Tensor, Tensor, Tensor]

torchgw._sinkhorn.adjoint_sinkhorn_vjp(T, reg, grad_T, tau_a=1.0, tau_b=1.0)[source]

dL/dC via implicit differentiation at the Sinkhorn fixed point.

Log-domain fixed-point conditions (tau = 1 balanced, tau < 1 KL-relaxed):

f = tau_a (log a - LSE_j(-C/ε + g)), g = tau_b (log b - LSE_i(-C/ε + f))

Their Jacobian is J = [[I, tau_a P], [tau_b R, I]] with P_{ij} = T_{ij}/ã_i and R_{ji} = T_{ij}/b̃_j, where ã = T·1 and b̃ = Tᵀ·1 are the row/column sums of T (equal to a, b at a balanced fixed point). The adjoint system Jᵀ·[λ_u, λ_v] = [r_u, r_v], r_u = (G ⊙ T)·1, r_v = (G ⊙ T)ᵀ·1, G = grad_T, is solved via the Schur complement S = I_K - tau_a tau_b Pᵀ Rᵀ.

Balanced case: S has a rank-1 null space spanned by b̃ (the potential shift f + c, g - c leaves T unchanged) and the RHS is orthogonal to 1_K, so adding 11ᵀ/K makes S nonsingular without changing the VJP. Relaxed case (tau_a tau_b < 1): S is already nonsingular.

Final VJP: dL/dC_{kl} = (T_{kl}/ε) · (-G_{kl} + tau_a λ_u_k/ã_k + tau_b λ_v_l/b̃_l)

Parameters:
Return type:

Tensor

GW loop

torchgw._loop.gw_loop(*, provider, T_init, step, fgw_alpha, C_lin, M, alpha, epsilon, lambda_ema_beta, max_iter, tol, min_iter_before_converge, differentiable, sink_dtype, verbose, verbose_every)[source]

Run the sampled GW iteration. Returns (T, err_list, n_iter, gw_cost).

provider : DistanceProvider or None (None = pure Wasserstein, Λ = C_lin) step : callable (Lambda, reg, verbose) -> T_new, see make_*_step

Parameters:
Return type:

tuple[Tensor, list, int, float]

torchgw._loop.gw_cost_estimate(provider, T, M, device, dtype)[source]

Monte Carlo estimate of the linearized GW cost Λ from M anchor pairs.

With anchors (k_m, l_m) ~ T and D_left[i, m] = d_X(i, k_m), D_tgt[j, m] = d_Y(j, l_m):

Λ ≈ mean_m(D_left²) − (2/M) D_left·D_tgtᵀ + mean_m(D_tgt²)

Only anchor-to-all distances are ever needed, so no N x N matrix is built.

Parameters:
Return type:

Tensor

Sampling

torchgw._sampling.sample_pairs_gpu(T, M)[source]

Sample M (row, col) pairs on GPU using torch.multinomial.

Avoids transferring the full (N, K) transport plan to CPU. Only transfers 2*M integers back.

Parameters:

  • T (Tensor of shape (N, K), non-negative, on any device)

  • M (number of pairs to sample)

Returns:

  • rows (ndarray of shape (M,))

  • cols (ndarray of shape (M,))

Return type:

tuple[ndarray, ndarray]