From 44d46149e7f12cf4f712332f70fad0a6f78a0299 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Vincent-Cuaz?= <cedvincentcuaz@gmail.com>
Date: Tue, 10 Sep 2024 23:46:02 +0200
Subject: [PATCH 1/8] merge

---
 RELEASES.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/RELEASES.md b/RELEASES.md
index cc18cc91b..277af7847 100644
--- a/RELEASES.md
+++ b/RELEASES.md
@@ -10,6 +10,8 @@
 - Improved `ot.plot.plot1D_mat` (PR #649)
 - Added `nx.det` (PR #649)
 - `nx.sqrtm` is now broadcastable (takes ..., d, d) inputs (PR #649)
+- restructure `ot.unbalanced` module (PR #658)
+- add `ot.unbalanced.lbfgsb_unbalanced2` and add flexible reference measure `c` in all unbalanced solvers (PR #658)
 
 #### Closed issues
 - Fixed `ot.gaussian` ignoring weights when computing means (PR #649, Issue #648)

From 78e09f690c2568334a718236158c5fddb9a60a9a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Vincent-Cuaz?= <cedvincentcuaz@gmail.com>
Date: Tue, 15 Sep 2026 17:47:40 +0200
Subject: [PATCH 2/8] complete tests

---
 RELEASES.md                   |   1 +
 ot/gromov/_quantized.py       | 374 +++++++++++++++++++++++-----------
 test/gromov/test_quantized.py |  71 +++++--
 3 files changed, 310 insertions(+), 136 deletions(-)

diff --git a/RELEASES.md b/RELEASES.md
index 82bd61421..753438877 100644
--- a/RELEASES.md
+++ b/RELEASES.md
@@ -19,6 +19,7 @@
 - Fix mean centering in `ot.dr.fda` and `ot.dr.wda`: `np.mean(X)` returned a scalar instead of the per-feature mean, so `proj` did not center the data as documented. In `ot.dr.fda` the same pattern in the class means made the between-class scatter matrix independent of which features separate the classes, and FDA returned a non-discriminant direction (PR #840)
 - `ot.dr.fda` and `ot.dr.wda` no longer modify the input array `X` in place (PR #840)
 - Fix `UnbalancedSinkhornTransport` `transform` failing with `AttributeError: 'NoneType' object has no attribute 'array_equal'` when `fit` was called with missing parameters (PR #837, Issue #650)
+- Fix quantized (F)GW solvers that ordered OT based on clusters and not initial node ordering (PR #857, Issue #786)
 
 ## 0.9.7.post1
 
diff --git a/ot/gromov/_quantized.py b/ot/gromov/_quantized.py
index 861cbfd46..d855bbcbb 100644
--- a/ot/gromov/_quantized.py
+++ b/ot/gromov/_quantized.py
@@ -40,6 +40,8 @@ def quantized_fused_gromov_wasserstein_partitioned(
     list_R2,
     list_p1,
     list_p2,
+    part1=None,
+    part2=None,
     MR=None,
     alpha=1.0,
     build_OT=False,
@@ -113,6 +115,12 @@ def quantized_fused_gromov_wasserstein_partitioned(
         List of node distributions within each partition of the source space.
     list_p : list of npart2 arrays,
         List of node distributions within each partition of the target space.
+    part1 : list of npart1 arrays, optional. Default is None.
+        List of arrays containing the indices of nodes in each partition of the source space.
+        Required as input if `build_OT=True`.
+    part2 : list of npart2 arrays, optional. Default is None.
+        List of arrays containing the indices of nodes in each partition of the target space.
+        Required as input if `build_OT=True`.
     MR : array-like, shape (npart1, npart2), optional. (Default is None)
         Metric cost matrix between features of representants across spaces.
     alpha: float, optional. Default is None.
@@ -156,19 +164,33 @@ def quantized_fused_gromov_wasserstein_partitioned(
 
     """
     if nx is None:
-        arr = [CR1, CR2, *list_R1, *list_R2, *list_p1, *list_p2]
-
-        if MR is not None:
-            arr.append(MR)
+        arr = [CR1, CR2, *list_R1, *list_R2, *list_p1, *list_p2, MR]
+        if build_OT:
+            arr += [*part1, *part2]
 
         nx = get_backend(*arr)
 
     npart1 = len(list_R1)
     npart2 = len(list_R2)
 
+    if (npart1 != len(list_p1)) or (npart2 != len(list_p2)):
+        raise ValueError(
+            f"""
+            Inconsistent number of partitions between list_R1 ({npart1}), list_p1 ({len(list_p1)}), 
+            and list_R2 ({npart2}), list_p2 ({len(list_p2)}).
+            """
+        )
+    if build_OT and ((npart1 != len(part1)) or (npart2 != len(part2))):
+        raise ValueError(
+            f"""
+            Inconsistent number of partitions in part1 ({len(part1)})
+            and part2 ({len(part2)}).
+            """
+        )
+
     # compute marginals for global alignment
-    pR1 = nx.from_numpy(list_to_array([nx.sum(p) for p in list_p1]))
-    pR2 = nx.from_numpy(list_to_array([nx.sum(q) for q in list_p2]))
+    pR1 = nx.from_numpy(list_to_array([nx.sum(p) for p in list_p1]), type_as=CR1)
+    pR2 = nx.from_numpy(list_to_array([nx.sum(q) for q in list_p2]), type_as=CR2)
 
     # compute global alignment
     if alpha == 1.0:
@@ -247,22 +269,25 @@ def quantized_fused_gromov_wasserstein_partitioned(
                     Ts_local[(i, j)] = res_1d
 
     if build_OT:
-        T_rows = []
-        for i in range(npart1):
-            list_Ti = []
-            for j in range(npart2):
-                if T_global[i, j] == 0.0:
-                    T_local = nx.zeros(
-                        (list_R1[i].shape[0], list_R2[j].shape[0]), type_as=T_global
-                    )
-                else:
-                    T_local = T_global[i, j] * Ts_local[(i, j)]
-                list_Ti.append(T_local)
-
-            Ti = nx.concatenate(list_Ti, axis=1)
-            T_rows.append(Ti)
-        T = nx.concatenate(T_rows, axis=0)
-
+        # Do memory-efficient alternatives by backend for mutable vs non-mutable tensors
+        if nx.__name__ in ("numpy", "torch"):
+            T = _build_full_transport_by_assignment(
+                T_global,
+                Ts_local,
+                part1,
+                part2,
+                nx,
+            )
+        else:
+            T = _build_full_transport_by_concatenation(
+                T_global,
+                Ts_local,
+                list_R1,
+                list_R2,
+                part1,
+                part2,
+                nx,
+            )
     else:
         T = None
 
@@ -273,6 +298,124 @@ def quantized_fused_gromov_wasserstein_partitioned(
         return T_global, Ts_local, T
 
 
+def _build_full_transport_by_assignment(
+    T_global,
+    Ts_local,
+    part1,
+    part2,
+    nx,
+):
+    """
+    Build the full transport matrix by assigning local transport blocks.
+
+    The local couplings are written directly into the rows and columns
+    identified by each source and target partition. This preserves the
+    original sample ordering without materializing a partition-ordered
+    intermediate matrix.
+
+    Parameters
+    ----------
+    T_global : array-like, shape (npart1, npart2)
+        Global transport between source and target partition representants.
+    Ts_local : dict
+        Dictionary of local transport matrices keyed by ``(i, j)``.
+    part1 : list of array-like
+        Source partition indices in the original sample ordering.
+    part2 : list of array-like
+        Target partition indices in the original sample ordering.
+    nx : backend
+        POT backend used to create and assign the transport matrix.
+
+    Returns
+    -------
+    T : array-like, shape (ns, nt)
+        Full transport matrix between the original source and target samples.
+    """
+    ns = sum(indices.shape[0] for indices in part1)
+    nt = sum(indices.shape[0] for indices in part2)
+
+    T = nx.zeros((ns, nt), type_as=T_global)
+
+    for i, rows in enumerate(part1):
+        for j, cols in enumerate(part2):
+            if T_global[i, j] != 0.0:
+                T_block = T_global[i, j] * Ts_local[(i, j)]
+                T[rows[:, None], cols[None, :]] = T_block
+
+    return T
+
+
+def _build_full_transport_by_concatenation(
+    T_global,
+    Ts_local,
+    list_R1,
+    list_R2,
+    part1,
+    part2,
+    nx,
+):
+    """
+    Build the full transport matrix by concatenating partition blocks.
+
+    Local couplings are first concatenated in source and target partition
+    order, then rows and columns are reordered according to the partition
+    indices so that the result uses the original sample ordering.
+
+    Parameters
+    ----------
+    T_global : array-like, shape (npart1, npart2)
+        Global transport between source and target partition representants.
+    Ts_local : dict
+        Dictionary of local transport matrices keyed by ``(i, j)``.
+    list_R1 : list of array-like
+        Source representative-to-sample relations, used to determine block
+        dimensions.
+    list_R2 : list of array-like
+        Target representative-to-sample relations, used to determine block
+        dimensions.
+    part1 : list of array-like
+        Source partition indices in the original sample ordering.
+    part2 : list of array-like
+        Target partition indices in the original sample ordering.
+    nx : backend
+        POT backend used to concatenate and reorder the transport matrix.
+
+    Returns
+    -------
+    T : array-like, shape (ns, nt)
+        Full transport matrix between the original source and target samples.
+    """
+    T_rows = []
+
+    for i in range(len(list_R1)):
+        T_blocks = []
+
+        for j in range(len(list_R2)):
+            if T_global[i, j] == 0.0:
+                T_block = nx.zeros(
+                    (list_R1[i].shape[0], list_R2[j].shape[0]),
+                    type_as=T_global,
+                )
+            else:
+                T_block = T_global[i, j] * Ts_local[(i, j)]
+
+            T_blocks.append(T_block)
+
+        T_rows.append(nx.concatenate(T_blocks, axis=1))
+
+    T = nx.concatenate(T_rows, axis=0)
+
+    perm1 = nx.concatenate(part1, axis=0)
+    perm2 = nx.concatenate(part2, axis=0)
+
+    T = T[nx.argsort(perm1)]
+    T = nx.transpose(T)
+    T = T[nx.argsort(perm2)]
+    T = nx.transpose(T)
+
+    return T
+
+
 def get_graph_partition(
     C, npart, part_method="random", F=None, alpha=1.0, random_state=0, nx=None
 ):
@@ -308,8 +451,8 @@ def get_graph_partition(
 
     Returns
     -------
-    part : array-like, shape (npart,)
-        Array of partition assignment for each node.
+    part : list of array-like, length npart
+        List of arrays containing the indices of nodes in each partition.
 
     References
     ----------
@@ -333,40 +476,38 @@ def get_graph_partition(
             stacklevel=2,
         )
 
-        part = np.arange(n)
+        part = list(nx.arange(n)[:, None])
 
     elif npart == 1:
-        part = np.zeros(n)
+        part = [nx.arange(n)]
 
     elif part_method == "random":
         # randomly partition the space
         random.seed(random_state)
-        part = list_to_array(random.choices(np.arange(npart), k=C.shape[0]))
-
-    elif part_method == "louvain":
-        C = nx.to_numpy(C0)
-        graph = from_numpy_array(C)
-        part_sets = louvain_communities(graph, seed=random_state)
-        part = np.zeros(n)
-        for iset_, set_ in enumerate(part_sets):
-            set_ = list(set_)
-            part[set_] = iset_
+        part_assignments = random.choices(np.arange(npart), k=n)
+        part = [
+            nx.from_numpy(np.where(np.array(part_assignments) == i)[0])
+            for i in range(npart)
+        ]
 
-    elif part_method == "fluid":
+    elif part_method in ["louvain", "fluid"]:
         C = nx.to_numpy(C0)
         graph = from_numpy_array(C)
-        part_sets = asyn_fluidc(graph, npart, seed=random_state)
-        part = np.zeros(n)
-        for iset_, set_ in enumerate(part_sets):
-            set_ = list(set_)
-            part[set_] = iset_
+        if part_method == "louvain":
+            part_sets = louvain_communities(graph, seed=random_state)
+        else:
+            part_sets = asyn_fluidc(graph, npart, seed=random_state)
+        part = [
+            nx.from_numpy(np.array(list(nodes)).astype(np.int64)) for nodes in part_sets
+        ]
 
     elif part_method == "spectral":
         C = nx.to_numpy(C0)
         sc = SpectralClustering(
             n_clusters=npart, random_state=random_state, affinity="precomputed"
         ).fit(C)
-        part = sc.labels_
+        labels = sc.labels_
+        part = [nx.from_numpy(np.where(labels == i)[0]) for i in range(npart)]
 
     elif part_method in ["GW", "FGW"]:
         raise ValueError(f"`part_method == {part_method}` not implemented yet.")
@@ -378,7 +519,8 @@ def get_graph_partition(
             {"random", "louvain", "fluid", "spectral", "GW", "FGW"}.
             """
         )
-    return nx.from_numpy(part, type_as=C0)
+
+    return part
 
 
 def get_graph_representants(C, part, rep_method="pagerank", random_state=0, nx=None):
@@ -391,8 +533,8 @@ def get_graph_representants(C, part, rep_method="pagerank", random_state=0, nx=N
     ----------
     C : array-like, shape (n, n)
         structure matrix.
-    part : array-like, shape (n,)
-        Array of partition assignment for each node.
+    part : list of array-like, length npart
+        List of arrays containing the indices of nodes in each partition.
     rep_method : str, optional. Default is 'pagerank'.
         Selection method for representant in each partition. Can be either 'random'
         i.e random sampling within each partition, or 'pagerank' to select a
@@ -404,9 +546,9 @@ def get_graph_representants(C, part, rep_method="pagerank", random_state=0, nx=N
 
     Returns
     -------
-    rep_indices : list, shape (npart,)
-        indices for representative node of each partition sorted
-        according to partition identifiers.
+    rep_indices : array-like, shape (npart,)
+        Array of indices for representative node of each partition sorted
+        according to partition order in `part` with same type as `C`.
 
     References
     ----------
@@ -415,34 +557,31 @@ def get_graph_representants(C, part, rep_method="pagerank", random_state=0, nx=N
 
     """
     if nx is None:
-        nx = get_backend(C, part)
+        nx = get_backend(C, *part)
 
     rep_indices = []
-    part_ids = nx.unique(part)
-    n_part_ids = part_ids.shape[0]
-    if n_part_ids == C.shape[0]:
-        rep_indices = nx.arange(n_part_ids)
+    n = C.shape[0]
+    n_part = len(part)
+    if n_part == n:
+        rep_indices = [indices[0] for indices in part]
 
     elif rep_method == "random":
         random.seed(random_state)
-        for id_, part_id in enumerate(part_ids):
-            indices = nx.where(part == part_id)[0]
+        for indices in part:
             rep_indices.append(random.choice(indices))
 
     elif rep_method == "pagerank":
         C0, part0 = C, part
         C = nx.to_numpy(C0)
-        part = nx.to_numpy(part0)
-        part_ids = np.unique(part)
+        part = [nx.to_numpy(indices) for indices in part]
 
-        for id_ in part_ids:
-            indices = np.where(part == id_)[0]
+        for indices in part:
             C_id = C[indices, :][:, indices]
             graph = from_numpy_array(C_id)
             pagerank_values = list(pagerank(graph).values())
             rep_idx = np.argmax(pagerank_values)
             rep_indices.append(indices[rep_idx])
-
+        C, part = C0, part0
     else:
         raise ValueError(
             f"""
@@ -450,7 +589,7 @@ def get_graph_representants(C, part, rep_method="pagerank", random_state=0, nx=N
             {"random", "pagerank"}.
             """
         )
-
+    rep_indices = nx.from_numpy(list_to_array(rep_indices).astype(np.int64))
     return rep_indices
 
 
@@ -471,11 +610,11 @@ def format_partitioned_graph(
         Structure matrix.
     p : array-like, shape (n,),
         Node distribution.
-    part : array-like, shape (n,)
-        Array of partition assignment for each node.
-    rep_indices : list of array-like of ints, shape (npart,)
-        indices for representative node of each partition sorted according to
-        partition identifiers.
+    part : list of array-like, length npart
+        List of arrays containing the indices of nodes in each partition.
+    rep_indices : array-like, shape (npart,)
+        Array of indices for representative node of each partition sorted
+        according to partition order in `part` with same type as `C`.
     F : array-like, shape (n, d), optional. (Default is None)
         Optional feature matrix aligned with the graph structure.
     M : array-like, shape (n, n), optional. (Default is None)
@@ -492,7 +631,7 @@ def format_partitioned_graph(
     CR : array-like, shape (npart, npart)
         Structure matrix between partition representants.
     list_R : list of npart arrays,
-        List of relations between a representant and nodes in its partition,
+        List of relations between a representant and the nodes in its partition,
         for each partition.
     list_p : list of npart arrays,
         List of node distributions within each partition.
@@ -506,13 +645,7 @@ def format_partitioned_graph(
 
     """
     if nx is None:
-        arr = [C, p, part]
-        if F is not None:
-            arr.append(F)
-        if M is not None:
-            arr.append(M)
-
-        nx = get_backend(*arr)
+        nx = get_backend(C, p, *part, rep_indices, F, M)
 
     if alpha != 1.0:
         if (M is None) or (F is None):
@@ -531,11 +664,8 @@ def format_partitioned_graph(
 
     list_R, list_p = [], []
 
-    part_ids = nx.unique(part)
-
-    for id_, part_id in enumerate(part_ids):
-        indices = nx.where(part == part_id)[0]
-        list_R.append(C_new[rep_indices[id_], indices])
+    for cluster_id, indices in enumerate(part):
+        list_R.append(C_new[rep_indices[cluster_id], indices])
         list_p.append(p[indices])
 
     if F is None:
@@ -727,14 +857,10 @@ def quantized_fused_gromov_wasserstein(
                 """
             )
 
-    arr = [C1, C2]
-    if C1_aux is not None:
-        arr.append(C1_aux)
-    else:
+    arr = [C1, C2, C1_aux, C2_aux, F1, F2]
+    if C1_aux is None:
         C1_aux = C1
-    if C2_aux is not None:
-        arr.append(C2_aux)
-    else:
+    if C2_aux is None:
         C2_aux = C2
     if p is not None:
         arr.append(list_to_array(p))
@@ -744,10 +870,6 @@ def quantized_fused_gromov_wasserstein(
         arr.append(list_to_array(q))
     else:
         q = unif(C2.shape[0], type_as=C1)
-    if F1 is not None:
-        arr.append(F1)
-    if F2 is not None:
-        arr.append(F1)
 
     nx = get_backend(*arr)
 
@@ -832,6 +954,8 @@ def quantized_fused_gromov_wasserstein(
         list_R2,
         list_p1,
         list_p2,
+        part1,
+        part2,
         MR,
         alpha,
         build_OT=True,
@@ -892,8 +1016,8 @@ def get_partition_and_representants_samples(
 
     Returns
     -------
-    part : array-like, shape (npart,)
-        Array of partition assignment for each node.
+    part : list of array-like, length npart
+        List of arrays containing the indices of nodes in each partition.
 
     rep_indices : list, shape (npart,)
         indices for representative node of each partition sorted
@@ -918,36 +1042,42 @@ def get_partition_and_representants_samples(
             stacklevel=2,
         )
 
-        part = nx.arange(n)
-        rep_indices = nx.arange(n)
+        part = list(nx.arange(n)[:, None])
+        rep_indices = [i for i in range(n)]
 
     elif npart == 1:
         random.seed(random_state)
-        part = nx.zeros(n)
-        rep_indices = [random.choice(nx.arange(n))]
+        part = [nx.arange(n)]
+        rep_indices = [random.choice(np.arange(n))]
 
     elif method == "random":
         # randomly partition the space
         random.seed(random_state)
-        part = list_to_array(random.choices(np.arange(npart), k=X.shape[0]))
-        part = nx.from_numpy(part, type_as=X0)
+        part_assignments = random.choices(np.arange(npart), k=X.shape[0])
+        part = [
+            nx.from_numpy(np.where(np.array(part_assignments) == i)[0])
+            for i in range(npart)
+        ]
 
         # randomly select representant in each partition
         rep_indices = []
-        part_ids = nx.unique(part)
-        for id_, part_id in enumerate(part_ids):
-            indices = nx.where(part == part_id)[0]
+        for indices_array in part:
+            indices = nx.to_numpy(indices_array)
             rep_indices.append(random.choice(indices))
 
     elif method == "kmeans":
         X = nx.to_numpy(X0)
         km = KMeans(n_clusters=npart, random_state=random_state).fit(X)
-        part = nx.from_numpy(km.labels_, type_as=X0)
+        labels = km.labels_
+        part = [
+            nx.from_numpy(np.where(labels == i)[0].astype(np.int64))
+            for i in range(npart)
+        ]
 
         rep_indices = []
-        for part_id in range(npart):
-            indices = nx.where(part == part_id)[0]
-            dists = dist(X[indices], km.cluster_centers_[part_id][None, :])
+        for i in range(npart):
+            indices = np.where(labels == i)[0]
+            dists = dist(X[indices], km.cluster_centers_[i][None, :])
             best_idx = indices[dists.argmin()]
             rep_indices.append(best_idx)
 
@@ -958,6 +1088,9 @@ def get_partition_and_representants_samples(
             """
         )
 
+    rep_indices = nx.from_numpy(list_to_array(rep_indices), type_as=part[0])
+    # print('part:', type(part), type(part[0]), part[0].dtype)
+    # print('rep_indices:', type(rep_indices), rep_indices.dtype)
     return part, rep_indices
 
 
@@ -976,11 +1109,11 @@ def format_partitioned_samples(X, p, part, rep_indices, F=None, alpha=1.0, nx=No
         Structure matrix.
     p : array-like, shape (n,),
         Node distribution.
-    part : array-like, shape (n,)
-        Array of partition assignment for each node.
-    rep_indices : list of array-like of ints, shape (npart,)
-        indices for representative node of each partition sorted according to
-        partition identifiers.
+    part : list of array-like, length npart
+        List of arrays containing the indices of nodes in each partition.
+    rep_indices : array-like, shape (npart,)
+        Array of indices for representative node of each partition sorted
+        according to partition identifiers.
     F : array-like, shape (n, p), optional. (Default is None)
         Optional feature matrix aligned with the samples.
     alpha: float, optional. Default is 1.
@@ -1009,7 +1142,7 @@ def format_partitioned_samples(X, p, part, rep_indices, F=None, alpha=1.0, nx=No
 
     """
     if nx is None:
-        arr = [X, p, part]
+        arr = [X, p, *part, rep_indices]
         if F is not None:
             arr.append(F)
 
@@ -1028,14 +1161,11 @@ def format_partitioned_samples(X, p, part, rep_indices, F=None, alpha=1.0, nx=No
 
     list_R, list_p = [], []
 
-    part_ids = nx.unique(part)
-
-    for id_, part_id in enumerate(part_ids):
-        indices = nx.where(part == part_id)[0]
-        structure_R = dist(X[indices], X[rep_indices[id_]][None, :])
+    for cluster_id, indices in enumerate(part):
+        structure_R = dist(X[indices], X[rep_indices[cluster_id]][None, :])
 
         if alpha != 1:
-            features_R = dist(F[indices], F[rep_indices[id_]][None, :])
+            features_R = dist(F[indices], F[rep_indices[cluster_id]][None, :])
         else:
             features_R = 0.0
 
@@ -1195,7 +1325,7 @@ def quantized_fused_gromov_wasserstein_samples(
             """
         )
 
-    arr = [X1, X2]
+    arr = [X1, X2, F1, F2]
     if p is not None:
         arr.append(list_to_array(p))
     else:
@@ -1204,10 +1334,6 @@ def quantized_fused_gromov_wasserstein_samples(
         arr.append(list_to_array(q))
     else:
         q = unif(X2.shape[0], type_as=X1)
-    if F1 is not None:
-        arr.append(F1)
-    if F2 is not None:
-        arr.append(F1)
 
     nx = get_backend(*arr)
 
@@ -1255,6 +1381,8 @@ def quantized_fused_gromov_wasserstein_samples(
         list_R2,
         list_p1,
         list_p2,
+        part1,
+        part2,
         MR,
         alpha,
         build_OT=True,
diff --git a/test/gromov/test_quantized.py b/test/gromov/test_quantized.py
index c3b80bb7d..5a14e9450 100644
--- a/test/gromov/test_quantized.py
+++ b/test/gromov/test_quantized.py
@@ -22,8 +22,10 @@ def test_quantized_gw(nx):
     C2 = rng.uniform(low=10.0, high=20.0, size=(n_samples, n_samples))
     C2 = (C2 + C2.T) / 2.0
 
-    p = ot.unif(n_samples)
-    q = ot.unif(n_samples)
+    p = np.arange(n_samples).astype(float)
+    p /= p.sum()
+    q = np.arange(n_samples).astype(float)
+    q /= q.sum()
 
     npart2 = 3
 
@@ -50,7 +52,7 @@ def test_quantized_gw(nx):
                 npart1,
                 npart2,
                 p,
-                None,
+                q,
                 C1,
                 None,
                 part_method=part_method,
@@ -63,7 +65,7 @@ def test_quantized_gw(nx):
                 C2b,
                 npart1,
                 npart2,
-                None,
+                pb,
                 qb,
                 None,
                 C2b,
@@ -80,6 +82,12 @@ def test_quantized_gw(nx):
                 T_globalb, Ts_localb, Tb = resb
 
             Tb = nx.to_numpy(Tb)
+            print("T.sum(0):", T.sum(0))
+            print("Tb.sum(0):", Tb.sum(0))
+            print("T.sum(1):", T.sum(1))
+            print("Tb.sum(1):", Tb.sum(1))
+            print("p:", p)
+            print("q:", q)
             # check constraints
             np.testing.assert_allclose(T, Tb, atol=1e-06)
             np.testing.assert_allclose(
@@ -114,7 +122,6 @@ def test_quantized_fgw(nx):
 
     p = ot.unif(n_samples)
     q = ot.unif(n_samples)
-
     npart1 = 2
     npart2 = 3
 
@@ -226,7 +233,17 @@ def test_quantized_fgw(nx):
     MRb = ot.dist(FR1b, FR2b)
 
     T_globalb, Ts_localb, _ = ot.gromov.quantized_fused_gromov_wasserstein_partitioned(
-        CR1b, CR2b, list_R1b, list_R2b, list_p1b, list_p2b, MRb, alpha, build_OT=False
+        CR1b,
+        CR2b,
+        list_R1b,
+        list_R2b,
+        list_p1b,
+        list_p2b,
+        None,
+        None,  # part useless when build_OT=False
+        MRb,
+        alpha,
+        build_OT=False,
     )
 
     T_globalb = nx.to_numpy(T_globalb)
@@ -264,7 +281,17 @@ def test_quantized_fgw(nx):
     # for non admissible values of alpha
     with pytest.raises(ValueError):
         ot.gromov.quantized_fused_gromov_wasserstein_partitioned(
-            CR1b, CR2b, list_R1b, list_R2b, list_p1b, list_p2b, MRb, 0, build_OT=False
+            CR1b,
+            CR2b,
+            list_R1b,
+            list_R2b,
+            list_p1b,
+            list_p2b,
+            None,
+            None,  # part useless when build_OT=False
+            MRb,
+            0,
+            build_OT=False,
         )
 
     # for non-consistent feature information provided
@@ -364,8 +391,8 @@ def test_quantized_fgw_samples(nx):
     F1 = rng.uniform(low=0.0, high=10, size=(n_samples_1, 3))
     F2 = rng.uniform(low=0.0, high=10, size=(n_samples_2, 3))
 
-    p = ot.unif(n_samples_1)
-    q = ot.unif(n_samples_2)
+    p = np.random.dirichlet(np.ones(n_samples_1))
+    q = np.random.dirichlet(np.ones(n_samples_2))
 
     npart1 = 2
     npart2 = 3
@@ -382,17 +409,18 @@ def test_quantized_fgw_samples(nx):
     for npart1 in [1, n_samples_1 + 1, 2]:
         log_tests = [True, False, True]
         count_mode = 0
-
+        print("--- npart:", npart1, "---")
         for method in methods:
+            print("method:", method, " nx:", nx.__name__)
             log_ = log_tests[count_mode]
             count_mode += 1
 
             res = ot.gromov.quantized_fused_gromov_wasserstein_samples(
-                X1, X2, npart1, npart2, p, None, F1, F2, alpha, method, log_
+                X1, X2, npart1, npart2, p, q, F1, F2, alpha, method, log_
             )
 
             resb = ot.gromov.quantized_fused_gromov_wasserstein_samples(
-                X1b, X2b, npart1, npart2, None, qb, F1b, F2b, alpha, method, log_
+                X1b, X2b, npart1, npart2, pb, qb, F1b, F2b, alpha, method, log_
             )
 
             if log_:
@@ -437,8 +465,25 @@ def test_quantized_fgw_samples(nx):
 
     MRb = ot.dist(FR1b, FR2b)
 
+    print("CR1b:", type(CR1b), CR1b.dtype)
+    print("CR2b", type(CR2b), CR2b.dtype)
+    print("list_R1b:", type(list_R1b), type(list_R1b[0]), list_R1b[0].dtype)
+    print("list_R2b:", type(list_R2b), type(list_R2b[0]), list_R2b[0].dtype)
+    print("list_p1b:", type(list_p1b), type(list_p1b[0]), list_p1b[0].dtype)
+    print("list_p2b:", type(list_p2b), type(list_p2b[0]), list_p2b[0].dtype)
+    print("MRb:", type(MRb), MRb.dtype)
     T_globalb, Ts_localb, _ = ot.gromov.quantized_fused_gromov_wasserstein_partitioned(
-        CR1b, CR2b, list_R1b, list_R2b, list_p1b, list_p2b, MRb, alpha, build_OT=False
+        CR1b,
+        CR2b,
+        list_R1b,
+        list_R2b,
+        list_p1b,
+        list_p2b,
+        None,
+        None,  # part useless when build_OT=False
+        MRb,
+        alpha,
+        build_OT=False,
     )
 
     T_globalb = nx.to_numpy(T_globalb)

From a489bf0260b89105202403c19036a6d7ed2df516 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Vincent-Cuaz?= <cedvincentcuaz@gmail.com>
Date: Thu, 17 Sep 2026 16:09:07 +0200
Subject: [PATCH 3/8] complete tests

---
 test/gromov/test_quantized.py | 44 +++++++++++++++++++++++++++++++----
 1 file changed, 40 insertions(+), 4 deletions(-)

diff --git a/test/gromov/test_quantized.py b/test/gromov/test_quantized.py
index 5a14e9450..937b9f238 100644
--- a/test/gromov/test_quantized.py
+++ b/test/gromov/test_quantized.py
@@ -203,7 +203,8 @@ def test_quantized_fgw(nx):
                 if key in logb.keys():
                     np.testing.assert_allclose(log[key], logb[key], atol=1e-06)
 
-    # complementary tests for utils functions
+    ### Complementary tests for utils functions
+    # checking consistency between wrapper and utils functions
     DF1b = ot.dist(F1b, F1b)
     DF2b = ot.dist(F2b, F2b)
     C1b_new = alpha * C1b + (1 - alpha) * DF1b
@@ -239,11 +240,11 @@ def test_quantized_fgw(nx):
         list_R2b,
         list_p1b,
         list_p2b,
-        None,
-        None,  # part useless when build_OT=False
+        part1b,
+        part2b,
         MRb,
         alpha,
-        build_OT=False,
+        build_OT=True,
     )
 
     T_globalb = nx.to_numpy(T_globalb)
@@ -312,6 +313,41 @@ def test_quantized_fgw(nx):
             "random",
             log_,
         )
+    ### Tests for non-consistent dimensions of inputs
+    # when build_OT = False, errors can come from inconsistent dimensions
+    # between list_R1b and list_p1b or list_R2b and list_p2b
+    with pytest.raises(ValueError):
+        ot.gromov.quantized_fused_gromov_wasserstein_partitioned(
+            CR1b,
+            CR2b,
+            list_R1b,
+            list_R2b,
+            list_p1b,
+            list_p2b[:-2],
+            None,
+            None,  # part useless when build_OT=False
+            MRb,
+            alpha,
+            build_OT=False,
+        )
+
+    # when build_OT = True, errors can also come from inconsistent dimensions
+    # between list_R1b and part1 or list_R2b and part2
+
+    with pytest.raises(ValueError):
+        ot.gromov.quantized_fused_gromov_wasserstein_partitioned(
+            CR1b,
+            CR2b,
+            list_R1b,
+            list_R2b,
+            list_p1b,
+            list_p2b,
+            part1b,
+            part2b[:-2],
+            MRb,
+            alpha,
+            build_OT=True,
+        )
 
 
 @pytest.skip_backend("jax", reason="test very slow with jax backend")

From acdb8f2e5b532a82e15f477f901d4f66396c58da Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Vincent-Cuaz?= <cedvincentcuaz@gmail.com>
Date: Fri, 18 Sep 2026 00:00:27 +0200
Subject: [PATCH 4/8] fix exemple and put notes in docstring for API change

---
 examples/gromov/plot_quantized_gromov_wasserstein.py |  2 ++
 ot/gromov/_quantized.py                              | 11 +++++++++++
 2 files changed, 13 insertions(+)

diff --git a/examples/gromov/plot_quantized_gromov_wasserstein.py b/examples/gromov/plot_quantized_gromov_wasserstein.py
index dcec838b6..a5b2171ac 100644
--- a/examples/gromov/plot_quantized_gromov_wasserstein.py
+++ b/examples/gromov/plot_quantized_gromov_wasserstein.py
@@ -259,6 +259,8 @@ def draw_graph(
     list_R2,
     list_h1,
     list_h2,
+    part1,
+    part2,
     MR=None,
     alpha=1.0,
     build_OT=True,
diff --git a/ot/gromov/_quantized.py b/ot/gromov/_quantized.py
index d855bbcbb..609b5cc01 100644
--- a/ot/gromov/_quantized.py
+++ b/ot/gromov/_quantized.py
@@ -100,6 +100,9 @@ def quantized_fused_gromov_wasserstein_partitioned(
         which can lead to copy overhead on GPU arrays.
     .. note:: All computations in the Gromov-Wasserstein conjugate gradient solver
         are done with numpy to limit memory overhead.
+    .. note:: The API changed in version 0.9.8 requesting `part1` and `part2`
+        as generated by `get_graph_partition` as inputs when `build_OT=True`
+        to fix inconsistencies in rows and columns ordering of the returned OT plan.
 
     Parameters
     ----------
@@ -425,6 +428,10 @@ def get_graph_partition(
     algorithms from networkx, or 'spectral' clustering from scikit-learn,
     or (Fused) Gromov-Wasserstein projections from POT.
 
+    .. note:: The API changed in version 0.9.8 outputting a list of node indices
+        per cluster instead of a single vector of node assignment to match the
+        new API of `quantized_fused_gromov_wasserstein_partitioned`.
+
     Parameters
     ----------
     C : array-like, shape (n, n)
@@ -997,6 +1004,10 @@ def get_partition_and_representants_samples(
     Compute `npart` partitions and representants over samples :math:`\mathbf{X} \in R^{n \times d}`
     using either a random or a kmeans algorithm.
 
+    .. note:: The API changed in version 0.9.8 outputting a list of node indices
+        per cluster in `part`, instead of a single vector of node assignment, to match the
+        new API of `quantized_fused_gromov_wasserstein_partitioned`.
+
     Parameters
     ----------
     X : array-like, shape (n, d)

From 1a8cd5c75a4ebf251cd21f1e153437e3b3e99038 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Vincent-Cuaz?= <cedvincentcuaz@gmail.com>
Date: Wed, 30 Sep 2026 16:20:15 +0200
Subject: [PATCH 5/8] fix exemple

---
 .../plot_quantized_gromov_wasserstein.py      | 83 ++++++++++++++-----
 1 file changed, 60 insertions(+), 23 deletions(-)

diff --git a/examples/gromov/plot_quantized_gromov_wasserstein.py b/examples/gromov/plot_quantized_gromov_wasserstein.py
index a5b2171ac..055fef20f 100644
--- a/examples/gromov/plot_quantized_gromov_wasserstein.py
+++ b/examples/gromov/plot_quantized_gromov_wasserstein.py
@@ -267,6 +267,24 @@ def draw_graph(
     log=True,
 )
 
+print("OT between representants: \n", OT_global_)
+print("OT local between samples and representants: \n", OTs_local_)
+
+for part_id1, cluster1 in enumerate(part1):
+    for part_id2, cluster2 in enumerate(part2):
+        print("sanity check (part1, part2) = (%s, %s): \n" % (part_id1, part_id2))
+        if OT_global_[part_id1, part_id2] > 0:
+            print(
+                np.all(
+                    OTs_local_[(part_id1, part_id2)] == OT_[cluster1, :][:, cluster2]
+                )
+            )
+plt.imshow(OT, interpolation="nearest", aspect="auto")
+plt.title("OT between samples")
+plt.colorbar()
+plt.axis("off")
+plt.show()
+
 
 # Visualization of the graph pre-processing
 
@@ -275,13 +293,17 @@ def draw_graph(
 seed_G1 = 0
 seed_G2 = 3
 
-part1_ = part1.astype(np.int32)
-part2_ = part2.astype(np.int32)
+part1_flat = np.zeros(C1.shape[0], dtype=np.int32)
+for cluster_id, cluster in enumerate(part1):
+    part1_flat[cluster] = cluster_id
 
+part2_flat = np.empty(C2.shape[0], dtype=np.int32)
+for cluster_id, cluster in enumerate(part2):
+    part2_flat[cluster] = cluster_id
 
-nodes_color_part1 = node_coloring(part1_, starting_color=0)
+nodes_color_part1 = node_coloring(part1_flat, starting_color=0)
 nodes_color_part2 = node_coloring(
-    part2_, starting_color=np.unique(nodes_color_part1).shape[0]
+    part2_flat, starting_color=np.unique(nodes_color_part1).shape[0]
 )
 
 
@@ -292,7 +314,7 @@ def draw_graph(
 pl.title(r"Input graph: $\mathbf{spC_1}$", fontsize=fontsize)
 
 pos1 = draw_graph(
-    G1, C1, ["C0" for _ in part1_], rep_indices1, node_size=node_size, seed=seed_G1
+    G1, C1, ["C0" for _ in part1_flat], rep_indices1, node_size=node_size, seed=seed_G1
 )
 
 pl.subplot(2, 3, 2)
@@ -320,7 +342,7 @@ def draw_graph(
 pl.title(r"Input graph: $\mathbf{spC_2}$", fontsize=fontsize)
 
 pos2 = draw_graph(
-    G2, C2, ["C0" for _ in part2_], rep_indices2, node_size=node_size, seed=seed_G2
+    G2, C2, ["C0" for _ in part2_flat], rep_indices2, node_size=node_size, seed=seed_G2
 )
 
 pl.subplot(2, 3, 5)
@@ -378,7 +400,14 @@ def draw_graph(
 
 qGW_dist = log["qFGW_dist"]
 
+print("OT between representants: \n", OT_global)
+print("OT local between samples and representants: \n", OTs_local)
 
+plt.imshow(OT, interpolation="nearest", aspect="auto")
+plt.title("OT between samples")
+plt.colorbar()
+plt.axis("off")
+plt.show()
 #############################################################################
 #
 # Visualization of the quantized Gromov-Wasserstein matching
@@ -394,8 +423,8 @@ def draw_transp_colored_qGW(
     C1,
     G2,
     C2,
-    part1,
-    part2,
+    part1_flat,
+    part2_flat,
     rep_indices1,
     rep_indices2,
     T,
@@ -410,15 +439,15 @@ def draw_transp_colored_qGW(
 ):
     starting_color = 0
     # get graphs partition and their coloring
-    unique_colors1 = ["C%s" % (starting_color + i) for i in np.unique(part1)]
+    unique_colors1 = ["C%s" % (starting_color + i) for i in np.unique(part1_flat)]
     nodes_color_part1 = []
-    for cluster in part1:
+    for cluster in part1_flat:
         nodes_color_part1.append(unique_colors1[cluster])
 
     starting_color = len(unique_colors1) + 1
-    unique_colors2 = ["C%s" % (starting_color + i) for i in np.unique(part2)]
+    unique_colors2 = ["C%s" % (starting_color + i) for i in np.unique(part2_flat)]
     nodes_color_part2 = []
-    for cluster in part2:
+    for cluster in part2_flat:
         nodes_color_part2.append(unique_colors2[cluster])
 
     pos1 = draw_graph(
@@ -440,7 +469,7 @@ def draw_transp_colored_qGW(
         pos=pos2,
         node_size=node_size,
         shiftx=shiftx,
-        seed=seed_G1,
+        seed=seed_G2,
         highlight_rep=highlight_rep,
     )
 
@@ -488,8 +517,8 @@ def draw_transp_colored_qGW(
     C1,
     weightedG2,
     C2,
-    part1_,
-    part2_,
+    part1_flat,
+    part2_flat,
     rep_indices1,
     rep_indices2,
     T=OT_,
@@ -512,8 +541,8 @@ def draw_transp_colored_qGW(
     C1,
     weightedG2,
     C2,
-    part1_,
-    part2_,
+    part1_flat,
+    part2_flat,
     rep_indices1,
     rep_indices2,
     T=OT_global,
@@ -568,8 +597,16 @@ def draw_transp_colored_qGW(
 part1, rep_indices1 = get_partition_and_representants_samples(X, 4, "kmeans", 0)
 part2, rep_indices2 = get_partition_and_representants_samples(Y, 4, "kmeans", 0)
 
-upart1 = np.unique(part1)
-upart2 = np.unique(part2)
+part1_flat = np.empty(X.shape[0], dtype=np.int32)
+for cluster_id, cluster in enumerate(part1):
+    part1_flat[cluster] = cluster_id
+
+part2_flat = np.empty(Y.shape[0], dtype=np.int32)
+for cluster_id, cluster in enumerate(part2):
+    part2_flat[cluster] = cluster_id
+
+upart1 = np.unique(part1_flat)
+upart2 = np.unique(part2_flat)
 
 # Plot the source and target samples as distributions
 s = 20
@@ -584,14 +621,14 @@ def draw_transp_colored_qGW(
 ax2 = fig.add_subplot(1, 3, 2)
 ax2.set_title("Partitioning")
 for i, elem in enumerate(upart1):
-    idx = np.argwhere(part1 == elem)[:, 0]
+    idx = np.where(part1_flat == elem)[0]
     ax2.scatter(X[idx, 0], X[idx, 1], color="C%s" % i, alpha=FX[idx], s=s)
 plt.axis("off")
 
 ax3 = fig.add_subplot(1, 3, 3)
 ax3.set_title("Representant selection")
 for i, elem in enumerate(upart1):
-    idx = np.argwhere(part1 == elem)[:, 0]
+    idx = np.where(part1_flat == elem)[0]
     ax3.scatter(X[idx, 0], X[idx, 1], color="C%s" % i, alpha=FX[idx], s=10)
     rep_idx = rep_indices1[i]
     ax3.scatter(
@@ -613,7 +650,7 @@ def draw_transp_colored_qGW(
 ax5 = fig.add_subplot(1, 3, 2, projection="3d")
 ax5.set_title("Partitioning")
 for i, elem in enumerate(upart2):
-    idx = np.argwhere(part2 == elem)[:, 0]
+    idx = np.where(part2_flat == elem)[0]
     color = "C%s" % (start_color + i)
     ax5.scatter(Y[idx, 0], Y[idx, 1], Y[idx, 2], c=color, alpha=FY[idx], s=s)
 plt.axis("off")
@@ -621,7 +658,7 @@ def draw_transp_colored_qGW(
 ax6 = fig.add_subplot(1, 3, 3, projection="3d")
 ax6.set_title("Representant selection")
 for i, elem in enumerate(upart2):
-    idx = np.argwhere(part2 == elem)[:, 0]
+    idx = np.where(part2_flat == elem)[0]
     color = "C%s" % (start_color + i)
     rep_idx = rep_indices2[i]
     ax6.scatter(Y[idx, 0], Y[idx, 1], Y[idx, 2], c=color, alpha=FY[idx], s=s)

From a8f75ac7be71440f95610fefb3475319ea39001a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Vincent-Cuaz?= <cedvincentcuaz@gmail.com>
Date: Sat, 3 Oct 2026 15:23:22 +0200
Subject: [PATCH 6/8] fix examples and tests

---
 .../plot_quantized_gromov_wasserstein.py      | 33 ++------
 ot/gromov/_quantized.py                       | 42 +++++-----
 test/gromov/test_quantized.py                 | 84 +++++++++++++++++++
 3 files changed, 115 insertions(+), 44 deletions(-)

diff --git a/examples/gromov/plot_quantized_gromov_wasserstein.py b/examples/gromov/plot_quantized_gromov_wasserstein.py
index 055fef20f..340229de9 100644
--- a/examples/gromov/plot_quantized_gromov_wasserstein.py
+++ b/examples/gromov/plot_quantized_gromov_wasserstein.py
@@ -216,7 +216,7 @@ def draw_graph(
 # to form partitioned graphs, which are then passed as input to the generic qFGW solver.
 
 # 1-a) Partition C1 and C2 in 2 and 3 clusters respectively using Louvain
-#    algorithm from Networkx. Then encode these partitions via vectors of assignments.
+#    algorithm from NetworkX. Then encode these partitions via vectors of assignments.
 
 part_method = "louvain"
 rep_method = "pagerank"
@@ -225,17 +225,16 @@ def draw_graph(
 npart_2 = 3  # 3 clusters used to describe C2
 
 part1 = get_graph_partition(
-    C1, npart=npart_1, part_method=part_method, F=None, alpha=1.0
+    C1, npart=npart_1, part_method=part_method, F=None, alpha=1.0, random_state=0
 )
 part2 = get_graph_partition(
-    C2, npart=npart_2, part_method=part_method, F=None, alpha=1.0
+    C2, npart=npart_2, part_method=part_method, F=None, alpha=1.0, random_state=0
 )
 
-# 1-b) Select representant in each partition using the Pagerank algorithm
-#     implementation from networkx.
+# 1-b) Select the PageRank representative in each partition.
 
-rep_indices1 = get_graph_representants(C1, part1, rep_method=rep_method)
-rep_indices2 = get_graph_representants(C2, part2, rep_method=rep_method)
+rep_indices1 = get_graph_representants(C1, part1, rep_method=rep_method, random_state=0)
+rep_indices2 = get_graph_representants(C2, part2, rep_method=rep_method, random_state=0)
 
 # 1-c) Format partitions such that:
 # CR contains relations between representants in each space.
@@ -267,19 +266,7 @@ def draw_graph(
     log=True,
 )
 
-print("OT between representants: \n", OT_global_)
-print("OT local between samples and representants: \n", OTs_local_)
-
-for part_id1, cluster1 in enumerate(part1):
-    for part_id2, cluster2 in enumerate(part2):
-        print("sanity check (part1, part2) = (%s, %s): \n" % (part_id1, part_id2))
-        if OT_global_[part_id1, part_id2] > 0:
-            print(
-                np.all(
-                    OTs_local_[(part_id1, part_id2)] == OT_[cluster1, :][:, cluster2]
-                )
-            )
-plt.imshow(OT, interpolation="nearest", aspect="auto")
+plt.imshow(OT_, interpolation="nearest", aspect="auto")
 plt.title("OT between samples")
 plt.colorbar()
 plt.axis("off")
@@ -395,14 +382,12 @@ def draw_graph(
     alpha=1.0,
     part_method=part_method,
     rep_method=rep_method,
+    random_state=0,
     log=True,
 )
 
 qGW_dist = log["qFGW_dist"]
 
-print("OT between representants: \n", OT_global)
-print("OT local between samples and representants: \n", OTs_local)
-
 plt.imshow(OT, interpolation="nearest", aspect="auto")
 plt.title("OT between samples")
 plt.colorbar()
@@ -521,7 +506,7 @@ def draw_transp_colored_qGW(
     part2_flat,
     rep_indices1,
     rep_indices2,
-    T=OT_,
+    T=OT,
     shiftx=1.5,
     node_size=node_size,
     seed_G1=seed_G1,
diff --git a/ot/gromov/_quantized.py b/ot/gromov/_quantized.py
index 609b5cc01..4047a8b22 100644
--- a/ot/gromov/_quantized.py
+++ b/ot/gromov/_quantized.py
@@ -830,30 +830,36 @@ def quantized_fused_gromov_wasserstein(
         Quantized gromov-wasserstein. ECML PKDD 2021. Springer International Publishing.
 
     """
-    if part_method in ["fluid", "louvain", "fluid_fused", "louvain_fused"] or (
-        rep_method in ["pagerank", "pagerank_fused"]
-    ):
-        if not networkx_import:
-            warnings.warn(
-                f"""
-                Networkx is not installed, so part_method={part_method} and/or
-                rep_method={rep_method} cannot be used and are set to `random`
-                default methods. Consider installing Networkx to fix this.
-                """
-            )
+    networkx_part_methods = {"fluid", "louvain", "fluid_fused", "louvain_fused"}
+    networkx_rep_methods = {"pagerank", "pagerank_fused"}
+    sklearn_part_methods = {"spectral", "spectral_fused"}
+
+    if not networkx_import and part_method in networkx_part_methods:
+        warnings.warn(
+            f"""
+            NetworkX is not installed, so part_method={part_method} cannot be
+            used and is set to `random`. Consider installing NetworkX to fix this.
+            """
+        )
         part_method = "random"
+
+    if not networkx_import and rep_method in networkx_rep_methods:
+        warnings.warn(
+            f"""
+            NetworkX is not installed, so rep_method={rep_method} cannot be
+            used and is set to `random`. Consider installing NetworkX to fix this.
+            """
+        )
         rep_method = "random"
 
-    if (part_method in ["spectral", "spectral_fused"]) and (not sklearn_import):
+    if not sklearn_import and part_method in sklearn_part_methods:
         warnings.warn(
             f"""
-            Scikit-learn is not installed, so part_method={part_method} and/or
-            rep_method={rep_method} cannot be used and are set to `random`
-            default methods. Consider installing Scikit-learn to fix this.
+            Scikit-learn is not installed, so part_method={part_method} cannot be
+            used and is set to `random`. Consider installing Scikit-learn to fix this.
             """
         )
         part_method = "random"
-        rep_method = "random"
 
     if ("fused" in part_method) or ("fused" in rep_method) or (part_method == "FGW"):
         if (F1 is None) or (F2 is None):
@@ -904,7 +910,6 @@ def quantized_fused_gromov_wasserstein(
         part2 = get_graph_partition(
             C2_aux, npart2, part_method, F2, alpha, random_state, nx
         )
-
     if "fused" in rep_method:
         if DF1 is None:
             DF1 = dist(F1, F1)
@@ -928,7 +933,6 @@ def quantized_fused_gromov_wasserstein(
         rep_indices2 = get_graph_representants(
             C2_aux, part2, rep_method, random_state, nx
         )
-
     # format partitions over (C1, F1) and (C2, F2)
     if (F1 is None) and (F2 is None):
         CR1, list_R1, list_p1 = format_partitioned_graph(
@@ -1100,8 +1104,6 @@ def get_partition_and_representants_samples(
         )
 
     rep_indices = nx.from_numpy(list_to_array(rep_indices), type_as=part[0])
-    # print('part:', type(part), type(part[0]), part[0].dtype)
-    # print('rep_indices:', type(rep_indices), rep_indices.dtype)
     return part, rep_indices
 
 
diff --git a/test/gromov/test_quantized.py b/test/gromov/test_quantized.py
index 937b9f238..6b6970514 100644
--- a/test/gromov/test_quantized.py
+++ b/test/gromov/test_quantized.py
@@ -4,14 +4,98 @@
 #
 # License: MIT License
 
+import warnings
+
 import numpy as np
 import pytest
 
 import ot
+from ot.gromov import _quantized
 
 from ot.gromov._quantized import networkx_import, sklearn_import
 
 
+@pytest.mark.parametrize(
+    "has_networkx,has_sklearn,requested_part,requested_rep,expected_part,expected_rep",
+    [
+        (True, True, "louvain", "pagerank", "louvain", "pagerank"),
+        (False, True, "louvain", "pagerank", "random", "random"),
+        (False, True, "random", "pagerank", "random", "random"),
+        (True, False, "spectral", "random", "random", "random"),
+        (True, False, "louvain", "pagerank", "louvain", "pagerank"),
+    ],
+)
+def test_quantized_gw_only_falls_back_for_missing_dependencies(
+    monkeypatch,
+    has_networkx,
+    has_sklearn,
+    requested_part,
+    requested_rep,
+    expected_part,
+    expected_rep,
+):
+    method_calls = {"part": [], "representant": []}
+
+    def fake_partition(
+        C, npart, part_method, F=None, alpha=1.0, random_state=0, nx=None
+    ):
+        method_calls["part"].append(part_method)
+        return [
+            np.asarray(indices)
+            for indices in np.array_split(np.arange(C.shape[0]), npart)
+        ]
+
+    def fake_representants(C, part, rep_method, random_state=0, nx=None):
+        method_calls["representant"].append(rep_method)
+        return np.asarray([indices[0] for indices in part])
+
+    def fake_partitioned_solver(
+        CR1, CR2, list_R1, list_R2, list_p1, list_p2, part1, part2, *args, **kwargs
+    ):
+        return (
+            np.zeros((len(list_R1), len(list_R2))),
+            {},
+            np.zeros((sum(map(len, part1)), sum(map(len, part2)))),
+        )
+
+    monkeypatch.setattr(_quantized, "networkx_import", has_networkx)
+    monkeypatch.setattr(_quantized, "sklearn_import", has_sklearn)
+    monkeypatch.setattr(_quantized, "get_graph_partition", fake_partition)
+    monkeypatch.setattr(_quantized, "get_graph_representants", fake_representants)
+    monkeypatch.setattr(
+        _quantized,
+        "quantized_fused_gromov_wasserstein_partitioned",
+        fake_partitioned_solver,
+    )
+
+    C1 = np.zeros((4, 4))
+    C2 = np.zeros((6, 6))
+    with warnings.catch_warnings(record=True) as recorded_warnings:
+        warnings.simplefilter("always")
+        _quantized.quantized_fused_gromov_wasserstein(
+            C1,
+            C2,
+            2,
+            3,
+            alpha=1.0,
+            part_method=requested_part,
+            rep_method=requested_rep,
+            random_state=0,
+        )
+
+    assert method_calls["part"] == [expected_part, expected_part]
+    assert method_calls["representant"] == [expected_rep, expected_rep]
+    expected_warning_count = sum(
+        [
+            not has_networkx
+            and requested_part in {"fluid", "louvain", "fluid_fused", "louvain_fused"},
+            not has_networkx and requested_rep in {"pagerank", "pagerank_fused"},
+            not has_sklearn and requested_part in {"spectral", "spectral_fused"},
+        ]
+    )
+    assert len(recorded_warnings) == expected_warning_count
+
+
 def test_quantized_gw(nx):
     n_samples = 30  # nb samples
 

From 0868de40481096451bdee345859d1caf9834eb3e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Vincent-Cuaz?= <cedvincentcuaz@gmail.com>
Date: Sat, 3 Oct 2026 15:33:21 +0200
Subject: [PATCH 7/8] remove forgotten prints

---
 test/gromov/test_quantized.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/test/gromov/test_quantized.py b/test/gromov/test_quantized.py
index 6b6970514..2cd6a61f3 100644
--- a/test/gromov/test_quantized.py
+++ b/test/gromov/test_quantized.py
@@ -585,13 +585,6 @@ def test_quantized_fgw_samples(nx):
 
     MRb = ot.dist(FR1b, FR2b)
 
-    print("CR1b:", type(CR1b), CR1b.dtype)
-    print("CR2b", type(CR2b), CR2b.dtype)
-    print("list_R1b:", type(list_R1b), type(list_R1b[0]), list_R1b[0].dtype)
-    print("list_R2b:", type(list_R2b), type(list_R2b[0]), list_R2b[0].dtype)
-    print("list_p1b:", type(list_p1b), type(list_p1b[0]), list_p1b[0].dtype)
-    print("list_p2b:", type(list_p2b), type(list_p2b[0]), list_p2b[0].dtype)
-    print("MRb:", type(MRb), MRb.dtype)
     T_globalb, Ts_localb, _ = ot.gromov.quantized_fused_gromov_wasserstein_partitioned(
         CR1b,
         CR2b,

From adc1d2b1728eca8d83d1b7d926b4cfa25414163b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Vincent-Cuaz?= <cedvincentcuaz@gmail.com>
Date: Sat, 3 Oct 2026 18:03:51 +0200
Subject: [PATCH 8/8] improve explanations in example + docstrings in functions

---
 .../plot_quantized_gromov_wasserstein.py      | 146 +++++++------
 ot/gromov/_quantized.py                       | 204 +++++++++---------
 2 files changed, 187 insertions(+), 163 deletions(-)

diff --git a/examples/gromov/plot_quantized_gromov_wasserstein.py b/examples/gromov/plot_quantized_gromov_wasserstein.py
index 340229de9..45bd9831f 100644
--- a/examples/gromov/plot_quantized_gromov_wasserstein.py
+++ b/examples/gromov/plot_quantized_gromov_wasserstein.py
@@ -1,39 +1,43 @@
 # -*- coding: utf-8 -*-
-"""
+r"""
 ===============================================
 Quantized Fused Gromov-Wasserstein examples
 ===============================================
 
 .. note::
-    Examples added in release: 0.9.4.
+    Examples added in release: 0.9.4 and updated with new API from 0.9.8.
 
 These examples show how to use the quantized (Fused) Gromov-Wasserstein
-solvers (qFGW) [68]. POT provides a generic solver `quantized_fused_gromov_wasserstein_partitioned`
+solvers (qFGW) [68]. POT provides a generic solver ``quantized_fused_gromov_wasserstein_partitioned``
 that takes as inputs partitioned graphs potentially endowed with node features,
 which have to be built by the user. On top of that, POT provides two wrappers:
-    i) `quantized_fused_gromov_wasserstein` operating over generic graphs, whose
-    partitioning is performed via `get_graph_partition` using e.g the Louvain algorithm,
-    and representant for each partition can be selected via `get_graph_representants`
+    i) ``quantized_fused_gromov_wasserstein`` operating over generic graphs, whose
+    partitioning is performed via ``get_graph_partition`` using e.g the Louvain algorithm,
+    and representant for each partition can be selected via ``get_graph_representants``
     using e.g the PageRank algorithm.
 
-    ii) `quantized_fused_gromov_wasserstein_samples` operating over point clouds,
-    e.g :math:`X_1 \in R^{n_1 * d_1}` and :math:`X_2 \in R^{n_2 * d_2}`
+    ii) ``quantized_fused_gromov_wasserstein_samples`` operating over point clouds,
+    e.g. :math:`X_1 \in \mathbb{R}^{n_1 \times d_1}` and :math:`X_2 \in \mathbb{R}^{n_2 \times d_2}`
     endowed with their respective euclidean geometry, whose partitioning and
     representant selection is performed jointly using e.g the K-means algorithm
-    via the function `get_partition_and_representants_samples`.
+    via the function ``get_partition_and_representants_samples``.
 
 
 We illustrate next how to compute the qGW distance on both types of data by:
 
-    i) Generating two graphs following Stochastic Block Models encoded as shortest
-    path matrices as qGW solvers tends to require dense structure to achieve a good
-    approximation of the GW distance (as qGW is an upper-bound of GW). In the meantime,
-    we illustrate an optional feature of our solvers, namely the use of auxiliary
-    structures e.g adjacency matrices to perform the graph partitioning.
+    i) Generating two graphs following stochastic block models. We represent each
+    graph by :math:`(\mathbf{C}_i, \mathbf{h}_i)`, where :math:`\mathbf{A}_i` is
+    its adjacency matrix, :math:`\mathbf{C}_i = d_{\mathrm{sp}}(\mathbf{A}_i)`
+    is its all-pairs shortest-path matrix, and :math:`\mathbf{h}_i` is the uniform
+    node distribution. qGW solvers tend to require dense structure to achieve a
+    good approximation of the GW distance (as qGW is an upper-bound of GW). We
+    also illustrate using :math:`\mathbf{A}_i` as an auxiliary structure for
+    graph partitioning.
 
     ii) Generating two point clouds representing curves in 2D and 3D respectively.
     We augment these point clouds by considering additional features of the same
-    dimensionaly :math:`F_1 \in R^{n_1 * d}` and :math:`F_2 \in R^{n_2 * d}`,
+    dimensionality :math:`\mathbf{F}_1 \in \mathbb{R}^{n_1 \times d}` and
+    :math:`\mathbf{F}_2 \in \mathbb{R}^{n_2 \times d}`,
     representing the color intensity associated to each sample of both distributions.
     Then we compute the qFGW distance between these attributed point clouds.
 
@@ -67,10 +71,12 @@
 
 #############################################################################
 #
-# Generate graphs
+# Generate graphs represented by (C1, h1) and (C2, h2)
 # --------------------------------------------------------------------------
 #
-# Create two graphs following Stochastic Block models of 2 and 3 clusters.
+# Create two stochastic block model graphs. Their adjacency matrices are A1 and
+# A2; their structure matrices C1 and C2 are the corresponding shortest-path
+# matrices, and h1 and h2 are uniform node distributions.
 
 N1 = 30  # 2 communities
 N2 = 45  # 3 communities
@@ -80,14 +86,14 @@
 G2 = sbm(seed=0, sizes=[N2 // 3, N2 // 3, N2 // 3], p=p2)
 
 
-C1 = networkx.to_numpy_array(G1)
-C2 = networkx.to_numpy_array(G2)
+A1 = networkx.to_numpy_array(G1)
+A2 = networkx.to_numpy_array(G2)
 
-spC1 = shortest_path(C1)
-spC2 = shortest_path(C2)
+C1 = shortest_path(A1)
+C2 = shortest_path(A2)
 
-h1 = np.ones(C1.shape[0]) / C1.shape[0]
-h2 = np.ones(C2.shape[0]) / C2.shape[0]
+h1 = np.ones(A1.shape[0]) / A1.shape[0]
+h2 = np.ones(A2.shape[0]) / A2.shape[0]
 
 # Add weights on the edges for visualization later on
 weight_intra_G1 = 5
@@ -215,8 +221,8 @@ def draw_graph(
 # We detail next the steps implemented within the wrapper that preprocess graphs
 # to form partitioned graphs, which are then passed as input to the generic qFGW solver.
 
-# 1-a) Partition C1 and C2 in 2 and 3 clusters respectively using Louvain
-#    algorithm from NetworkX. Then encode these partitions via vectors of assignments.
+# 1-a) Partition the graphs represented by (C1, h1) and (C2, h2), using their
+#    adjacency matrices A1 and A2 as inputs to the Louvain algorithm.
 
 part_method = "louvain"
 rep_method = "pagerank"
@@ -225,16 +231,16 @@ def draw_graph(
 npart_2 = 3  # 3 clusters used to describe C2
 
 part1 = get_graph_partition(
-    C1, npart=npart_1, part_method=part_method, F=None, alpha=1.0, random_state=0
+    A1, npart=npart_1, part_method=part_method, F=None, alpha=1.0, random_state=0
 )
 part2 = get_graph_partition(
-    C2, npart=npart_2, part_method=part_method, F=None, alpha=1.0, random_state=0
+    A2, npart=npart_2, part_method=part_method, F=None, alpha=1.0, random_state=0
 )
 
 # 1-b) Select the PageRank representative in each partition.
 
-rep_indices1 = get_graph_representants(C1, part1, rep_method=rep_method, random_state=0)
-rep_indices2 = get_graph_representants(C2, part2, rep_method=rep_method, random_state=0)
+rep_indices1 = get_graph_representants(A1, part1, rep_method=rep_method, random_state=0)
+rep_indices2 = get_graph_representants(A2, part2, rep_method=rep_method, random_state=0)
 
 # 1-c) Format partitions such that:
 # CR contains relations between representants in each space.
@@ -242,11 +248,11 @@ def draw_graph(
 # list_h contains samples relative importance within each partition.
 
 CR1, list_R1, list_h1 = format_partitioned_graph(
-    spC1, h1, part1, rep_indices1, F=None, M=None, alpha=1.0
+    C1, h1, part1, rep_indices1, F=None, M=None, alpha=1.0
 )
 
 CR2, list_R2, list_h2 = format_partitioned_graph(
-    spC2, h2, part2, rep_indices2, F=None, M=None, alpha=1.0
+    C2, h2, part2, rep_indices2, F=None, M=None, alpha=1.0
 )
 
 # 1-d) call to partitioned quantized gromov-wasserstein solver
@@ -280,11 +286,11 @@ def draw_graph(
 seed_G1 = 0
 seed_G2 = 3
 
-part1_flat = np.zeros(C1.shape[0], dtype=np.int32)
+part1_flat = np.zeros(A1.shape[0], dtype=np.int32)
 for cluster_id, cluster in enumerate(part1):
     part1_flat[cluster] = cluster_id
 
-part2_flat = np.empty(C2.shape[0], dtype=np.int32)
+part2_flat = np.empty(A2.shape[0], dtype=np.int32)
 for cluster_id, cluster in enumerate(part2):
     part2_flat[cluster] = cluster_id
 
@@ -298,17 +304,19 @@ def draw_graph(
 pl.clf()
 pl.axis("off")
 pl.subplot(2, 3, 1)
-pl.title(r"Input graph: $\mathbf{spC_1}$", fontsize=fontsize)
+pl.title(
+    r"Input graph: $\mathbf{C_1}=d_{\mathrm{sp}}(\mathbf{A}_1)$", fontsize=fontsize
+)
 
 pos1 = draw_graph(
-    G1, C1, ["C0" for _ in part1_flat], rep_indices1, node_size=node_size, seed=seed_G1
+    G1, A1, ["C0" for _ in part1_flat], rep_indices1, node_size=node_size, seed=seed_G1
 )
 
 pl.subplot(2, 3, 2)
 pl.title("Partitioning", fontsize=fontsize)
 
 _ = draw_graph(
-    G1, C1, nodes_color_part1, rep_indices1, pos=pos1, node_size=node_size, seed=seed_G1
+    G1, A1, nodes_color_part1, rep_indices1, pos=pos1, node_size=node_size, seed=seed_G1
 )
 
 pl.subplot(2, 3, 3)
@@ -316,7 +324,7 @@ def draw_graph(
 
 _ = draw_graph(
     G1,
-    C1,
+    A1,
     nodes_color_part1,
     rep_indices1,
     pos=pos1,
@@ -326,17 +334,19 @@ def draw_graph(
 )
 
 pl.subplot(2, 3, 4)
-pl.title(r"Input graph: $\mathbf{spC_2}$", fontsize=fontsize)
+pl.title(
+    r"Input graph: $\mathbf{C_2}=d_{\mathrm{sp}}(\mathbf{A}_2)$", fontsize=fontsize
+)
 
 pos2 = draw_graph(
-    G2, C2, ["C0" for _ in part2_flat], rep_indices2, node_size=node_size, seed=seed_G2
+    G2, A2, ["C0" for _ in part2_flat], rep_indices2, node_size=node_size, seed=seed_G2
 )
 
 pl.subplot(2, 3, 5)
 pl.title(r"Partitioning", fontsize=fontsize)
 
 _ = draw_graph(
-    G2, C2, nodes_color_part2, rep_indices2, pos=pos2, node_size=node_size, seed=seed_G2
+    G2, A2, nodes_color_part2, rep_indices2, pos=pos2, node_size=node_size, seed=seed_G2
 )
 
 pl.subplot(2, 3, 6)
@@ -344,7 +354,7 @@ def draw_graph(
 
 _ = draw_graph(
     G2,
-    C2,
+    A2,
     nodes_color_part2,
     rep_indices2,
     pos=pos2,
@@ -359,24 +369,23 @@ def draw_graph(
 # Compute the quantized Gromov-Wasserstein distance using the wrapper
 # ---------------------------------------------------------
 #
-# Compute qGW(spC1, h1, spC2, h2). We also illustrate the use of auxiliary matrices
-# such that the adjacency matrices `C1_aux=C1` and `C2_aux=C2` to partition the graph using
-# Louvain algorithm, and the Pagerank algorithm for selecting representant within
-# each partition. Notice that `C1_aux` and `C2_aux` are optional, if they are not
-# specified these pre-processing algorithms will be applied to spC2 and spC3.
+# Compute qGW between the graph representations (C1, h1) and (C2, h2). The
+# adjacency matrices A1 and A2 are passed as auxiliary structures for Louvain
+# partitioning and PageRank representative selection; the wrapper uses C1 and C2
+# as the shortest-path structures in the transport problem.
 
 
 # no node features are considered on this synthetic dataset. Hence we simply
 # let F1, F2 = None and set alpha = 1.
 OT_global, OTs_local, OT, log = quantized_fused_gromov_wasserstein(
-    spC1,
-    spC2,
+    C1,
+    C2,
     npart_1,
     npart_2,
     h1,
     h2,
-    C1_aux=C1,
-    C2_aux=C2,
+    C1_aux=A1,
+    C2_aux=A2,
     F1=None,
     F2=None,
     alpha=1.0,
@@ -405,9 +414,9 @@ def draw_graph(
 
 def draw_transp_colored_qGW(
     G1,
-    C1,
+    A1,
     G2,
-    C2,
+    A2,
     part1_flat,
     part2_flat,
     rep_indices1,
@@ -437,7 +446,7 @@ def draw_transp_colored_qGW(
 
     pos1 = draw_graph(
         G1,
-        C1,
+        A1,
         nodes_color_part1,
         rep_indices1,
         pos=pos1,
@@ -448,7 +457,7 @@ def draw_transp_colored_qGW(
     )
     pos2 = draw_graph(
         G2,
-        C2,
+        A2,
         nodes_color_part2,
         rep_indices2,
         pos=pos2,
@@ -493,15 +502,15 @@ def draw_transp_colored_qGW(
 pl.axis("off")
 pl.subplot(1, 2, 1)
 pl.title(
-    r"qGW$(\mathbf{spC_1}, \mathbf{spC_1}) =%s$" % (np.round(qGW_dist, 3)),
+    r"qGW$(\mathbf{C_1}, \mathbf{C_2}) =%s$" % (np.round(qGW_dist, 3)),
     fontsize=fontsize,
 )
 
 pos1, pos2 = draw_transp_colored_qGW(
     weightedG1,
-    C1,
+    A1,
     weightedG2,
-    C2,
+    A2,
     part1_flat,
     part2_flat,
     rep_indices1,
@@ -523,9 +532,9 @@ def draw_transp_colored_qGW(
 
 pos1, pos2 = draw_transp_colored_qGW(
     weightedG1,
-    C1,
+    A1,
     weightedG2,
-    C2,
+    A2,
     part1_flat,
     part2_flat,
     rep_indices1,
@@ -546,8 +555,13 @@ def draw_transp_colored_qGW(
 # Generate attributed point clouds
 # --------------------------------------------------------------------------
 #
-# Create two attributed point clouds representing curves in 2D and 3D respectively,
-# whose samples are further associated to various color intensities.
+# Create point clouds :math:`\mathbf{X}_1 \in \mathbb{R}^{n \times 2}` and
+# :math:`\mathbf{X}_2 \in \mathbb{R}^{n \times 3}`, with feature matrices
+# :math:`\mathbf{F}_1, \mathbf{F}_2 \in \mathbb{R}^{n \times 1}`. The attributed
+# point-cloud inputs are :math:`(\mathbf{D}(\mathbf{X}_i), \mathbf{F}_i,
+# \mathbf{h}_i)`, where :math:`\mathbf{D}(\mathbf{X}_i)` is the pairwise
+# squared-Euclidean distance matrix and :math:`\mathbf{h}_i` is the uniform
+# distribution over points.
 
 n_samples = 100
 
@@ -665,10 +679,12 @@ def draw_transp_colored_qGW(
 # Compute the quantized Fused Gromov-Wasserstein distance between samples using the wrapper
 # ---------------------------------------------------------
 #
-# Compute qFGW(X, FX, hX, Y, FY, HY), setting the trade-off parameter between
-# structures and features `alpha=0.5`. This solver considers a squared euclidean structure
-# for each distribution X and Y, and partition each of them into 4 clusters using
-# the K-means algorithm before computing qFGW.
+# Compute qFGW between :math:`(\mathbf{D}(\mathbf{X}_1), \mathbf{F}_1,
+# \mathbf{h}_1)` and :math:`(\mathbf{D}(\mathbf{X}_2), \mathbf{F}_2,
+# \mathbf{h}_2)`, with structure-feature trade-off :math:`\alpha=0.5`. In the
+# variables below, X1=X, X2=Y, F1=FX, and F2=FY; the omitted p and q arguments
+# make h1 and h2 uniform. The wrapper uses squared-Euclidean structure matrices
+# and partitions both point clouds into four clusters using K-means.
 
 T_global, Ts_local, T, log = quantized_fused_gromov_wasserstein_samples(
     X,
diff --git a/ot/gromov/_quantized.py b/ot/gromov/_quantized.py
index 4047a8b22..13a9bd5d0 100644
--- a/ot/gromov/_quantized.py
+++ b/ot/gromov/_quantized.py
@@ -60,24 +60,26 @@ def quantized_fused_gromov_wasserstein_partitioned(
     :math:`\mathcal{P_1} = \{(\mathbf{P_{1, i}}, \mathbf{r_{1, i}})\}_{i \leq npart1}`
     and :math:`\mathcal{P_2} = \{(\mathbf{P_{2, j}}, \mathbf{r_{2, j}})\}_{j \leq npart2}`.
     The latter must be precomputed and encoded e.g for the source as: :math:`\mathbf{CR_1}`
-    structure matrix between representants; `list_R1` a list of relations between
-    representants and their associated samples; `list_p1` a list of nodes
+    structure matrix between representants; ``list_R1`` a list of relations between
+    representants and their associated samples; ``list_p1`` a list of nodes
     distribution within each partition; :math:`\mathbf{FR_1}` feature matrix
     of representants.
 
     The function estimates the following optimization problem:
 
     .. math::
-        \mathbf{T}^* \in \mathop{\arg \min}_\mathbf{T} \quad \alpha \sum_{i,j,k,l}
-        L(\mathbf{C_1}_{i,k}, \mathbf{C_2}_{j,l}) \mathbf{T}_{i,j} \mathbf{T}_{k,l}
-        + (1-\alpha) \langle \mathbf{T}, M\rangle_F
-        s.t. \ \mathbf{T} \mathbf{1} &= \mathbf{p}
-
-             \mathbf{T}^T \mathbf{1} &= \mathbf{q}
-
-             \mathbf{T} &\geq 0
-
-             \mathbf{T}_{|\mathbf{P_{1, i}}, \mathbf{P_{2, j}}} &= T^{g}_{ij} \mathbf{T}^{(i,j)}
+                \begin{aligned}
+                \mathbf{T}^* \in \mathop{\arg\min}_{\mathbf{T}} \quad
+                & \alpha \sum_{i,j,k,l} L(\mathbf{C_1}_{i,k}, \mathbf{C_2}_{j,l})
+                    \mathbf{T}_{i,j} \mathbf{T}_{k,l} \\
+                & + (1-\alpha) \langle \mathbf{T}, M\rangle_F \\
+                \mathrm{s.t.} \quad
+                & \mathbf{T} \mathbf{1} = \mathbf{p}, \\
+                & \mathbf{T}^T \mathbf{1} = \mathbf{q}, \\
+                & \mathbf{T} \geq 0, \\
+                & \mathbf{T}_{|\mathbf{P_{1, i}}, \mathbf{P_{2, j}}}
+                    = T^{g}_{ij} \mathbf{T}^{(i,j)}.
+                \end{aligned}
 
     using a two-step strategy computing: i) a global alignment :math:`\mathbf{T}^{g}`
     between representants joint structure and feature spaces; ii) local alignments
@@ -100,8 +102,8 @@ def quantized_fused_gromov_wasserstein_partitioned(
         which can lead to copy overhead on GPU arrays.
     .. note:: All computations in the Gromov-Wasserstein conjugate gradient solver
         are done with numpy to limit memory overhead.
-    .. note:: The API changed in version 0.9.8 requesting `part1` and `part2`
-        as generated by `get_graph_partition` as inputs when `build_OT=True`
+    .. note:: The API changed in version 0.9.8 requesting ``part1`` and ``part2``
+        as generated by ``get_graph_partition`` as inputs when ``build_OT=True``
         to fix inconsistencies in rows and columns ordering of the returned OT plan.
 
     Parameters
@@ -120,15 +122,15 @@ def quantized_fused_gromov_wasserstein_partitioned(
         List of node distributions within each partition of the target space.
     part1 : list of npart1 arrays, optional. Default is None.
         List of arrays containing the indices of nodes in each partition of the source space.
-        Required as input if `build_OT=True`.
+        Required as input if ``build_OT=True``.
     part2 : list of npart2 arrays, optional. Default is None.
         List of arrays containing the indices of nodes in each partition of the target space.
-        Required as input if `build_OT=True`.
+        Required as input if ``build_OT=True``.
     MR : array-like, shape (npart1, npart2), optional. (Default is None)
         Metric cost matrix between features of representants across spaces.
     alpha: float, optional. Default is None.
         FGW trade-off parameter in :math:`]0, 1]` between structure and features.
-        If `alpha = 1` features are ignored hence computing qGW.
+        If ``alpha = 1`` features are ignored hence computing qGW.
     build_OT: bool, optional. Default is False
         Either to build or not the OT between non-partitioned structures.
     log : bool, optional. Default is False
@@ -146,18 +148,18 @@ def quantized_fused_gromov_wasserstein_partitioned(
         POT backend
 
     **kwargs : dict
-        parameters can be directly passed to the ot.optim.cg solver
+        parameters can be directly passed to the ``ot.optim.cg`` solver
 
     Returns
     -------
-    T_global: array-like, shape (`npart1`, `npart2`)
+    T_global: array-like, shape (``npart1``, ``npart2``)
         Gromov-Wasserstein alignment :math:`\mathbf{T}^{g}` between representants.
     Ts_local: dict of local OT matrices.
-        Dictionary with keys :math:`(i, j)` corresponding to 1D OT between
+        Dictionary with keys :math:`(i, j)` corresponding to one-dimensional OT between
         :math:`\mathbf{P_{1, i}}` and :math:`\mathbf{P_{2, j}}` if :math:`T^{g}_{ij} \neq 0`.
-    T: array-like, shape `(ns, nt)`
-        Coupling between the two spaces if `build_OT=True` else None.
-    log : dict, if `log=True`.
+    T: array-like, shape ``(ns, nt)``
+        Coupling between the two spaces if ``build_OT=True`` else None.
+    log : dict, if ``log=True``.
         Convergence information and losses of inner OT problems.
 
     References
@@ -239,7 +241,7 @@ def quantized_fused_gromov_wasserstein_partitioned(
     else:
         raise ValueError(
             f"""
-            `alpha='{alpha}'` should be in ]0, 1].
+            ``alpha='{alpha}'`` should be in ]0, 1].
             """
         )
 
@@ -424,13 +426,13 @@ def get_graph_partition(
 ):
     r"""
     Partitioning a given graph with structure matrix :math:`\mathbf{C} \in R^{n \times n}`
-    into `npart` partitions either 'random', or using one of {'louvain', 'fluid'}
-    algorithms from networkx, or 'spectral' clustering from scikit-learn,
+    into ``npart`` partitions either ``random``, or using one of {``louvain``, ``fluid``}
+    algorithms from NetworkX, or ``spectral`` clustering from scikit-learn,
     or (Fused) Gromov-Wasserstein projections from POT.
 
     .. note:: The API changed in version 0.9.8 outputting a list of node indices
         per cluster instead of a single vector of node assignment to match the
-        new API of `quantized_fused_gromov_wasserstein_partitioned`.
+        new API of ``quantized_fused_gromov_wasserstein_partitioned``.
 
     Parameters
     ----------
@@ -439,18 +441,18 @@ def get_graph_partition(
     npart : int,
         number of partitions/clusters smaller than the number of nodes in
         :math:`\mathbf{C}`.
-    part_method : str, optional. Default is 'random'.
-        Partitioning algorithm to use among {'random', 'louvain', 'fluid', 'spectral', 'GW', 'FGW'}.
-        'random' for random sampling of points; 'louvain' and 'fluid' for graph
-        partitioning algorithm that works well on adjacency matrix, If the
-        louvain algorithm is used, `npart` is ignored; 'spectral' for spectral
-        clustering; '(F)GW' for (F)GW projection using sr(F)GW solvers.
+    part_method : str, optional. Default is ``random``.
+        Partitioning algorithm to use among {``random``, ``louvain``, ``fluid``, ``spectral``, ``GW``, ``FGW``}.
+        ``random`` for random sampling of points; ``louvain`` and ``fluid`` for graph
+        partitioning algorithms that work well on adjacency matrices. If
+        ``louvain`` is used, ``npart`` is ignored; ``spectral`` denotes spectral
+        clustering; ``GW`` and ``FGW`` denote (F)GW projection using sr(F)GW solvers.
     F : array-like, shape (n, d), optional. (Default is None)
         Optional feature matrix aligned with the graph structure. Only used if
-        `part_method="FGW"`.
+        ``part_method="FGW"``.
     alpha : float, optional. (Default is 1.)
         Trade-off parameter between feature and structure matrices, taking
-        values in [0, 1] and only used if `F != None` and `part_method="FGW"`.
+        values in [0, 1] and only used if ``F != None`` and ``part_method="FGW"``.
     random_state: int, optional
         Random seed for the partitioning algorithm.
     nx : backend, optional
@@ -534,7 +536,7 @@ def get_graph_representants(C, part, rep_method="pagerank", random_state=0, nx=N
     r"""
     Get representative node for each partition given by :math:`\mathbf{part} \in R^{n}`
     of a graph with structure matrix :math:`\mathbf{C} \in R^{n \times n}`.
-    Selection is either done randomly or using 'pagerank' algorithm from networkx.
+    Selection is either done randomly or using the ``pagerank`` algorithm from NetworkX.
 
     Parameters
     ----------
@@ -542,9 +544,9 @@ def get_graph_representants(C, part, rep_method="pagerank", random_state=0, nx=N
         structure matrix.
     part : list of array-like, length npart
         List of arrays containing the indices of nodes in each partition.
-    rep_method : str, optional. Default is 'pagerank'.
-        Selection method for representant in each partition. Can be either 'random'
-        i.e random sampling within each partition, or 'pagerank' to select a
+    rep_method : str, optional. Default is ``pagerank``.
+        Selection method for the representative in each partition. Can be ``random``
+        (random sampling within each partition), or ``pagerank`` to select a
         node with maximal pagerank.
     random_state: int, optional
         Random seed for the partitioning algorithm
@@ -555,7 +557,7 @@ def get_graph_representants(C, part, rep_method="pagerank", random_state=0, nx=N
     -------
     rep_indices : array-like, shape (npart,)
         Array of indices for representative node of each partition sorted
-        according to partition order in `part` with same type as `C`.
+        according to partition order in ``part`` with same type as ``C``.
 
     References
     ----------
@@ -621,14 +623,14 @@ def format_partitioned_graph(
         List of arrays containing the indices of nodes in each partition.
     rep_indices : array-like, shape (npart,)
         Array of indices for representative node of each partition sorted
-        according to partition order in `part` with same type as `C`.
+        according to partition order in ``part`` with same type as ``C``.
     F : array-like, shape (n, d), optional. (Default is None)
         Optional feature matrix aligned with the graph structure.
     M : array-like, shape (n, n), optional. (Default is None)
         Optional pairwise similarity matrix between features.
     alpha: float, optional. Default is 1.
         Trade-off parameter in :math:`]0, 1]` between structure and features.
-        If `alpha = 1` features are ignored. This trade-off is taken into account
+        If ``alpha = 1`` features are ignored. This trade-off is taken into account
         into the outputted relations between nodes and representants.
     nx : backend, optional
         POT backend
@@ -642,7 +644,7 @@ def format_partitioned_graph(
         for each partition.
     list_p : list of npart arrays,
         List of node distributions within each partition.
-    FR : array-like, shape (npart, d), if `F != None`.
+    FR : array-like, shape (npart, d), if ``F != None``.
         Feature matrix of representants.
 
     References
@@ -715,16 +717,19 @@ def quantized_fused_gromov_wasserstein(
     The function estimates the following optimization problem:
 
     .. math::
-        \mathbf{T}^* \in \mathop{\arg \min}_\mathbf{T} \quad \alpha \sum_{i,j,k,l}
-        L(\mathbf{C_1}_{i,k}, \mathbf{C_2}_{j,l}) \mathbf{T}_{i,j} \mathbf{T}_{k,l}
-        + (1-\alpha) \langle \mathbf{T}, \mathbf{D}(\mathbf{F_1}, \mathbf{F}_2) \rangle_F
-        s.t. \ \mathbf{T} \mathbf{1} &= \mathbf{p}
-
-             \mathbf{T}^T \mathbf{1} &= \mathbf{q}
-
-             \mathbf{T} &\geq 0
-
-             \mathbf{T}_{|\mathbf{P_{1, i}}, \mathbf{P_{2, j}}} &= T^{g}_{ij} \mathbf{T}^{(i,j)}
+                \begin{aligned}
+                \mathbf{T}^* \in \mathop{\arg\min}_{\mathbf{T}} \quad
+                & \alpha \sum_{i,j,k,l} L(\mathbf{C_1}_{i,k}, \mathbf{C_2}_{j,l})
+                    \mathbf{T}_{i,j} \mathbf{T}_{k,l} \\
+                & + (1-\alpha) \langle \mathbf{T},
+                    \mathbf{D}(\mathbf{F_1}, \mathbf{F}_2) \rangle_F \\
+                \mathrm{s.t.} \quad
+                & \mathbf{T} \mathbf{1} = \mathbf{p}, \\
+                & \mathbf{T}^T \mathbf{1} = \mathbf{q}, \\
+                & \mathbf{T} \geq 0, \\
+                & \mathbf{T}_{|\mathbf{P_{1, i}}, \mathbf{P_{2, j}}}
+                    = T^{g}_{ij} \mathbf{T}^{(i,j)}.
+                \end{aligned}
 
     using a two-step strategy computing: i) a global alignment :math:`\mathbf{T}^{g}`
     between representants across joint structure and feature spaces;
@@ -776,23 +781,23 @@ def quantized_fused_gromov_wasserstein(
         Feature matrix in the target space
     alpha: float, optional. Default is 1.
         FGW trade-off parameter in :math:`]0, 1]` between structure and features.
-        If `alpha = 1` features are ignored hence computing qGW, if `alpha=0`
+        If ``alpha = 1`` features are ignored hence computing qGW, if ``alpha=0``
         structures are ignored and we compute the quantized Wasserstein transport.
-    part_method : str, optional. Default is 'spectral'.
-        Partitioning algorithm to use among {'random', 'louvain', 'fluid',
-        'spectral', 'louvain_fused', 'fluid_fused', 'spectral_fused', 'GW', 'FGW'}.
-        If part_method in {'louvain_fused', 'fluid_fused', 'spectral_fused'},
-        corresponding graph partitioning algorithm {'louvain', 'fluid', 'spectral'}
+    part_method : str, optional. Default is ``spectral``.
+        Partitioning algorithm to use among {``random``, ``louvain``, ``fluid``,
+        ``spectral``, ``louvain_fused``, ``fluid_fused``, ``spectral_fused``, ``GW``, ``FGW``}.
+        If ``part_method`` is one of {``louvain_fused``, ``fluid_fused``, ``spectral_fused``},
+        the corresponding graph partitioning algorithm (``louvain``, ``fluid``, or ``spectral``)
         will be used on the modified structure matrix
         :math:`\alpha \mathbf{C} + (1 - \alpha) \mathbf{D}(\mathbf{F})` where
         :math:`\mathbf{D}(\mathbf{F})` is the pairwise euclidean matrix between features.
-        If part_method in {'GW', 'FGW'}, a (F)GW projection is used.
-        If the louvain algorithm is used, the requested number of partitions is
+        If ``part_method`` is ``GW`` or ``FGW``, a (F)GW projection is used.
+        If ``louvain`` is used, the requested number of partitions is
         ignored.
-    rep_method : str, optional. Default is 'pagerank'.
+    rep_method : str, optional. Default is ``pagerank``.
         Selection method for node representant in each partition.
-        Can be either 'random' i.e random sampling within each partition,
-        {'pagerank', 'pagerank_fused'} to select a node with maximal pagerank w.r.t
+        Can be either ``random`` (random sampling within each partition), or
+        ``pagerank`` or ``pagerank_fused`` to select a node with maximal pagerank with respect to
         :math:`\mathbf{C}` or :math:`\alpha \mathbf{C} + (1 - \alpha) \mathbf{D}(\mathbf{F})`.
     verbose : bool, optional
         Print information along iterations
@@ -814,12 +819,12 @@ def quantized_fused_gromov_wasserstein(
 
     Returns
     -------
-    T_global: array-like, shape (`npart1`, `npart2`)
+    T_global: array-like, shape (``npart1``, ``npart2``)
         Fused Gromov-Wasserstein alignment :math:`\mathbf{T}^{g}` between representants.
     Ts_local: dict of local OT matrices.
         Dictionary with keys :math:`(i, j)` corresponding to 1D OT between
         :math:`\mathbf{P_{1, i}}` and :math:`\mathbf{P_{2, j}}` if :math:`T^{g}_{ij} \neq 0`.
-    T: array-like, shape `(ns, nt)`
+    T: array-like, shape ``(ns, nt)``
         Coupling between the two spaces.
     log : dict
         Convergence information for inner problems and qGW loss.
@@ -865,7 +870,7 @@ def quantized_fused_gromov_wasserstein(
         if (F1 is None) or (F2 is None):
             raise ValueError(
                 f"""
-                `part_method='{part_method}'` and/or `rep_method='{rep_method}'`
+                ``part_method='{part_method}'`` and/or ``rep_method='{rep_method}'``
                 require feature matrices which are not provided as inputs.
                 """
             )
@@ -1005,12 +1010,12 @@ def get_partition_and_representants_samples(
     X, npart, method="kmeans", random_state=0, nx=None
 ):
     r"""
-    Compute `npart` partitions and representants over samples :math:`\mathbf{X} \in R^{n \times d}`
-    using either a random or a kmeans algorithm.
+    Compute ``npart`` partitions and representatives over samples :math:`\mathbf{X} \in R^{n \times d}`
+    using either ``random`` partitioning or ``kmeans`` clustering.
 
     .. note:: The API changed in version 0.9.8 outputting a list of node indices
-        per cluster in `part`, instead of a single vector of node assignment, to match the
-        new API of `quantized_fused_gromov_wasserstein_partitioned`.
+        per cluster in ``part``, instead of a single vector of node assignment, to match the
+        new API of ``quantized_fused_gromov_wasserstein_partitioned``.
 
     Parameters
     ----------
@@ -1019,11 +1024,10 @@ def get_partition_and_representants_samples(
     npart : int,
         number of partitions smaller than the number of samples in
         :math:`\mathbf{X}`.
-    method : str, optional. Default is 'kmeans'.
-        Partitioning and representant selection algorithms to use among
-        {'random', 'kmeans'}. 'random' for random sampling of points; 'kmeans'
-        for k-means clustering using scikit-learn implementation where closest
-        points to centroids are considered as representants.
+    method : str, optional. Default is ``kmeans``.
+        Partitioning and representative-selection algorithms: ``random`` for
+        random partitioning and representatives, or ``kmeans`` to use scikit-learn
+        clustering with the points closest to cluster centers as representatives.
     random_state: int, optional
         Random seed for the partitioning algorithm.
     nx : backend, optional
@@ -1099,7 +1103,7 @@ def get_partition_and_representants_samples(
     else:
         raise ValueError(
             f"""
-            Unknown `method='{method}'`. Use one of: {"random", "kmeans"}
+            Unknown method '{method}'. Use either 'random' or 'kmeans'.
             """
         )
 
@@ -1127,11 +1131,11 @@ def format_partitioned_samples(X, p, part, rep_indices, F=None, alpha=1.0, nx=No
     rep_indices : array-like, shape (npart,)
         Array of indices for representative node of each partition sorted
         according to partition identifiers.
-    F : array-like, shape (n, p), optional. (Default is None)
+    F : array-like, shape (n, d), optional. (Default is None)
         Optional feature matrix aligned with the samples.
     alpha: float, optional. Default is 1.
         Trade-off parameter in :math:`]0, 1]` between structure and features.
-        If `alpha = 1` features are ignored. This trade-off is taken into account
+        If ``alpha = 1`` features are ignored. This trade-off is taken into account
         into the outputted relations between nodes and representants.
     nx : backend, optional
         POT backend
@@ -1145,7 +1149,7 @@ def format_partitioned_samples(X, p, part, rep_indices, F=None, alpha=1.0, nx=No
         for each partition.
     list_p : list of npart arrays,
         List of node distributions within each partition.
-    FR : array-like, shape (npart, d), if `F != None`.
+    FR : array-like, shape (npart, d), if ``F != None``.
         Feature matrix of representants.
 
     References
@@ -1214,24 +1218,28 @@ def quantized_fused_gromov_wasserstein_samples(
 ):
     r"""
     Returns the quantized Fused Gromov-Wasserstein transport between samples
-    endowed with their respective euclidean geometry :math:`(\mathbf{D}(\mathbf{X_1}), \mathbf{F_1}, \mathbf{p})`
-    and :math:`(\mathbf{D}(\mathbf{X_1}), \mathbf{F_2}, \mathbf{q})`, whose samples are assigned to partitions and
+    endowed with their respective Euclidean geometries :math:`(\mathbf{D}(\mathbf{X_1}), \mathbf{F_1}, \mathbf{p})`
+    and :math:`(\mathbf{D}(\mathbf{X_2}), \mathbf{F_2}, \mathbf{q})`, whose samples are assigned to partitions and
     representants :math:`\mathcal{P_1} = \{(\mathbf{P_{1, i}}, \mathbf{r_{1, i}})\}_{i \leq npart1}`
     and :math:`\mathcal{P_2} = \{(\mathbf{P_{2, j}}, \mathbf{r_{2, j}})\}_{j \leq npart2}`.
 
     The function estimates the following optimization problem:
 
     .. math::
-        \mathbf{T}^* \in \mathop{\arg \min}_\mathbf{T} \quad \alpha \sum_{i,j,k,l}
-        L(\mathbf{D}(\mathbf{X_1})_{i,k}, \mathbf{D}(\mathbf{X_2})_{j,l}) \mathbf{T}_{i,j} \mathbf{T}_{k,l}
-        + (1-\alpha) \langle \mathbf{T}, \mathbf{D}(\mathbf{F_1}, \mathbf{F}_2) \rangle_F
-        s.t. \ \mathbf{T} \mathbf{1} &= \mathbf{p}
-
-             \mathbf{T}^T \mathbf{1} &= \mathbf{q}
-
-             \mathbf{T} &\geq 0
-
-             \mathbf{T}_{|\mathbf{P_{1, i}}, \mathbf{P_{2, j}}} &= T^{g}_{ij} \mathbf{T}^{(i,j)}
+                \begin{aligned}
+                \mathbf{T}^* \in \mathop{\arg\min}_{\mathbf{T}} \quad
+                & \alpha \sum_{i,j,k,l}
+                    L(\mathbf{D}(\mathbf{X_1})_{i,k}, \mathbf{D}(\mathbf{X_2})_{j,l})
+                    \mathbf{T}_{i,j} \mathbf{T}_{k,l} \\
+                & + (1-\alpha) \langle \mathbf{T},
+                    \mathbf{D}(\mathbf{F_1}, \mathbf{F_2}) \rangle_F \\
+                \mathrm{s.t.} \quad
+                & \mathbf{T} \mathbf{1} = \mathbf{p}, \\
+                & \mathbf{T}^T \mathbf{1} = \mathbf{q}, \\
+                & \mathbf{T} \geq 0, \\
+                & \mathbf{T}_{|\mathbf{P_{1, i}}, \mathbf{P_{2, j}}}
+                    = T^{g}_{ij} \mathbf{T}^{(i,j)}.
+                \end{aligned}
 
     using a two-step strategy computing: i) a global alignment :math:`\mathbf{T}^{g}`
     between representants across joint structure and feature spaces;
@@ -1277,12 +1285,12 @@ def quantized_fused_gromov_wasserstein_samples(
         Feature matrix in the target space
     alpha: float, optional. Default is 1.
         FGW trade-off parameter in :math:`]0, 1]` between structure and features.
-        If `alpha = 1` features are ignored hence computing qGW, if `alpha=0`
+        If ``alpha = 1`` features are ignored hence computing qGW; if ``alpha=0``
         structures are ignored and we compute the quantized Wasserstein transport.
-    method : str, optional. Default is 'kmeans'.
+    method : str, optional. Default is ``kmeans``.
         Partitioning and representant selection algorithms to use among
-        {'random', 'kmeans', 'kmeans_fused'}.
-        If `part_method == 'kmeans_fused'`, kmeans is performed on augmented
+        {``random``, ``kmeans``, ``kmeans_fused``}.
+        If ``method`` is ``kmeans_fused``, k-means is performed on augmented
         samples :math:`[\alpha \mathbf{X}; (1 - \alpha) \mathbf{F}]`.
     verbose : bool, optional
         Print information along iterations
@@ -1304,12 +1312,12 @@ def quantized_fused_gromov_wasserstein_samples(
 
     Returns
     -------
-    T_global: array-like, shape (`npart1`, `npart2`)
+    T_global: array-like, shape (``npart1``, ``npart2``)
         Fused Gromov-Wasserstein alignment :math:`\mathbf{T}^{g}` between representants.
     Ts_local: dict of local OT matrices.
         Dictionary with keys :math:`(i, j)` corresponding to 1D OT between
         :math:`\mathbf{P_{1, i}}` and :math:`\mathbf{P_{2, j}}` if :math:`T^{g}_{ij} \neq 0`.
-    T: array-like, shape `(ns, nt)`
+    T: array-like, shape ``(ns, nt)``
         Coupling between the two spaces.
     log : dict
         Convergence information for inner problems and qGW loss.