fix: deep review

README.md

@@ -30,8 +30,8 @@ diarization output.
 
 ```sh
 # Pinned upstream revision + expected SHA-256 of the FP32 single-file ONNX.
-DIA_EMBED_MODEL_REV="38168b544a562dec24d49e63786c16e80782eeaf"
-DIA_EMBED_MODEL_SHA256="4c15c6be4235318d092c9d347e00c68ba476136d6172f675f76ad6b0c2661f01"
+DIA_EMBED_MODEL_REV="6eef479c954ec180e79cee316af2f16d5f7720bd"
+DIA_EMBED_MODEL_SHA256="f23f04aa9d0f6b8b0a28de016d226dcbe92d7461a6e58045401acfbed623838a"
 mkdir -p models
 TMP="$(mktemp "${TMPDIR:-/tmp}/wespeaker_resnet34_lm.XXXXXXXXXX")"
 ```

examples/run_owned_pipeline.rs

@@ -13,8 +13,11 @@
 //! to compute DER vs pyannote.
 
 use diarization::{
-  embed::EmbedModel, offline::OwnedDiarizationPipeline, plda::PldaTransform,
-  reconstruct::spans_to_rttm_lines, segment::SegmentModel,
+  embed::EmbedModel,
+  offline::{OwnedDiarizationPipeline, OwnedPipelineOptions},
+  plda::PldaTransform,
+  reconstruct::spans_to_rttm_lines,
+  segment::SegmentModel,
 };
 use std::path::PathBuf;
 
@@ -58,7 +61,13 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     .map_err(|e| format!("load embed model from {}: {}", emb_path.display(), e))?;
   let plda = PldaTransform::new()?;
 
-  let pipeline = OwnedDiarizationPipeline::new();
+  // `OwnedPipelineOptions::new()` defaults to `smoothing_epsilon =
+  // None` for bit-exact pyannote community-1 RTTM. Callers wanting
+  // speakrs-style streaming-friendly stable speaker assignments
+  // (sub-100ms overlap-region splits merged into the previously-
+  // selected speaker) opt in via `with_smoothing_epsilon(Some(eps))`.
+  let opts = OwnedPipelineOptions::new();
+  let pipeline = OwnedDiarizationPipeline::with_options(opts);
   let out = pipeline.run(&mut seg, &mut emb, &plda, &samples)?;
 
   // Use clip basename as the RTTM uri.

scripts/download-embed-model.sh

@@ -29,15 +29,15 @@ mkdir -p "$MODELS_DIR"
 # Pin a specific HF commit so the download is reproducible. The
 # README quickstart pins the same revision + SHA-256 inline; keep
 # both in sync when bumping.
-REV="38168b544a562dec24d49e63786c16e80782eeaf"
+REV="6eef479c954ec180e79cee316af2f16d5f7720bd"
 URL="https://huggingface.co/FinDIT-Studio/dia-models/resolve/$REV/wespeaker_resnet34_lm.onnx"
 DEST="$MODELS_DIR/wespeaker_resnet34_lm.onnx"
 
 # SHA-256 of the canonical packed FP32 model (single-file, no
 # external data) at the pinned `$REV`. Update both if the upstream
 # HF repo re-publishes — a mismatch indicates content drift that
 # could silently invalidate byte-determinism / pyannote-parity gates.
-EXPECTED_SHA256="4c15c6be4235318d092c9d347e00c68ba476136d6172f675f76ad6b0c2661f01"
+EXPECTED_SHA256="f23f04aa9d0f6b8b0a28de016d226dcbe92d7461a6e58045401acfbed623838a"
 
 if [ -f "$DEST" ]; then
   ACTUAL_SHA256="$(shasum -a 256 "$DEST" | awk '{print $1}')"

scripts/fix_wespeaker_pooling_eps.py

@@ -0,0 +1,136 @@
+"""Patch the WeSpeaker ONNX export to match pyannote's PyTorch
+statistics pooling on sparse-mask edge cases.
+
+Pyannote's `pyannote.audio.models.blocks.pooling.StatsPool` (line
+52, 58) computes weighted mean/std via:
+
+    v1 = weights.sum(dim=2) + 1e-8                     # eps for mean
+    mean = (sequences * weights).sum(dim=2) / v1
+    v2 = (weights ** 2).sum(dim=2)
+    var = ((seq - mean)**2 * weights).sum(dim=2) / (v1 - v2/v1 + 1e-8)
+    std = sqrt(var)
+
+The ONNX export shipped under `models/wespeaker_resnet34_lm.onnx`
+omits both `+ 1e-8` epsilons. With binary masks that have only 1-2
+active frames out of 589, this causes:
+
+  - 1 active frame:  v1 = 1, v2 = 1 → v1 - v2/v1 = 0 → div-by-zero → +inf
+                     → propagates through Gemm to f32::MAX-class
+                     embedding corruption (we measured 10/964 (chunk,
+                     speaker) pairs on testaudioset 10 with this).
+  - 2 active frames: v1 = 2, v2 = 2 → denom = 1, but f32 cancellation
+                     in `v1 - v2/v1` near edge can still amplify.
+
+The patch inserts two `Add(small_eps)` nodes:
+  - `sum_1_eps = sum_1 + 1e-8`  (used by both mean and var denoms)
+  - `sub_349_eps = sub_349 + 1e-8`  (used by var denom)
+
+Output `models/wespeaker_resnet34_lm_stable.onnx` matches pyannote's
+PyTorch stats pooling bit-exact for any mask sparsity.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import onnx
+from onnx import helper, numpy_helper, TensorProto
+
+EPS = 1e-8
+
+
+def patch(in_path: Path, out_path: Path) -> None:
+    m = onnx.load(str(in_path))
+    g = m.graph
+
+    # Add a 1e-8 constant initializer.
+    eps_init = numpy_helper.from_array(
+        np.array(EPS, dtype=np.float32), name="stats_pool_eps"
+    )
+    g.initializer.append(eps_init)
+
+    # Find the relevant tensor names by walking nodes.
+    # Node "ReduceSum" producing sum_1 (the v1 = sum(weights) tensor).
+    # Node names from the captured graph dump:
+    #   102: ReduceSum(unsqueeze_2) → sum_1
+    #   105: Div(sum_2, sum_1) → div  (this is the MEAN)
+    #   113: Div(sum_3, sum_1) → div_1
+    #   114: Sub(sum_1, div_1) → sub_349  (n_eff = v1 - v2/v1)
+    #   115: Div(sum_4, sub_349) → div_2  (this is var)
+    # We need:
+    #   sum_1     → sum_1 + eps  (for ALL consumers: 105 and 113)
+    #   sub_349   → sub_349 + eps  (for consumer: 115)
+    sum_1_consumers = ["div", "div_1"]      # nodes 105, 113 take sum_1
+    sub_349_consumers = ["div_2"]           # node 115 takes sub_349
+
+    # New tensor names.
+    sum_1_eps_name = "sum_1_eps"
+    sub_349_eps_name = "sub_349_eps"
+
+    # Insert Add nodes.
+    add_sum1 = helper.make_node(
+        "Add",
+        inputs=["sum_1", eps_init.name],
+        outputs=[sum_1_eps_name],
+        name="add_sum1_eps",
+    )
+    add_sub349 = helper.make_node(
+        "Add",
+        inputs=["sub_349", eps_init.name],
+        outputs=[sub_349_eps_name],
+        name="add_sub349_eps",
+    )
+
+    # Insert before any consumer that's a Div node. ONNX is
+    # topologically ordered, so insert right after the original
+    # producer. We append at the end and let ONNX reorder; in practice
+    # all our target consumers come AFTER the producers, so simple
+    # append works.
+    g.node.append(add_sum1)
+    g.node.append(add_sub349)
+
+    # Re-route consumers' inputs.
+    for n in g.node:
+        for i, inp in enumerate(n.input):
+            if inp == "sum_1" and n.output and n.output[0] in sum_1_consumers:
+                # Mean (node 105) and div_1 (node 113) both consume
+                # sum_1; pyannote uses v1 (sum_1+eps) for both.
+                n.input[i] = sum_1_eps_name
+            elif inp == "sub_349" and n.output and n.output[0] in sub_349_consumers:
+                # Variance denominator (node 115) — gets +eps.
+                n.input[i] = sub_349_eps_name
+
+    # Re-topologically-sort: Add nodes come right after their producer
+    # so consumers (which appear later in the original order) can see
+    # the new tensors. We rebuild the node list by pulling Add nodes
+    # forward into the right position.
+    nodes = list(g.node)
+    # Find positions.
+    sum_1_idx = next(i for i, n in enumerate(nodes) if n.output and n.output[0] == "sum_1")
+    sub_349_idx = next(i for i, n in enumerate(nodes) if n.output and n.output[0] == "sub_349")
+    # Remove the appended Add nodes from the end.
+    nodes = [n for n in nodes if n.name not in {"add_sum1_eps", "add_sub349_eps"}]
+    # Insert after their producers (later index first to keep earlier index stable).
+    insert_first = max(sum_1_idx, sub_349_idx)
+    insert_second = min(sum_1_idx, sub_349_idx)
+    if sum_1_idx > sub_349_idx:
+        nodes.insert(insert_first + 1, add_sum1)
+        nodes.insert(insert_second + 1, add_sub349)
+    else:
+        nodes.insert(insert_first + 1, add_sub349)
+        nodes.insert(insert_second + 1, add_sum1)
+    # Rebuild graph.
+    del g.node[:]
+    g.node.extend(nodes)
+
+    onnx.checker.check_model(m)
+    onnx.save(m, str(out_path))
+    print(f"[patch] {in_path.name} -> {out_path.name}: added 2 Add(+1e-8) nodes")
+
+
+if __name__ == "__main__":
+    in_p = Path(sys.argv[1] if len(sys.argv) > 1 else "models/wespeaker_resnet34_lm.onnx")
+    out_p = Path(sys.argv[2] if len(sys.argv) > 2 else "models/wespeaker_resnet34_lm_stable.onnx")
+    patch(in_p, out_p)

src/cluster/ahc/algo.rs

@@ -223,22 +223,17 @@ fn l2_normalize_to_row_major(
 /// downstream clustering correctness (the labels are arbitrary
 /// integers naming the buckets; DER is invariant to relabeling).
 ///
-/// **TODO**: if a future end-to-end parity test runs
-/// `ahc_init → build qinit → vbx_iterate → q_final` and compares
-/// element-wise against captured `q_final`, the `qinit` column ordering
-/// will not match (since our labels are a permutation of scipy's). At
-/// that point, choose one of:
-/// 1. Implement scipy's exact tree-traversal label order here (drop
-///    this canonicalization pass; align DFS push order with scipy's
-///    `_hierarchy.pyx::cluster_dist`).
-/// 2. Compare `q_final` modulo column permutation (mathematically
-///    equivalent — the permutation is recoverable from
-///    `(our_labels, scipy_labels)` matching).
-/// 3. Have `ahc_init` return `(labels, permutation_to_scipy)` so the
-///    caller can build the column-permuted qinit explicitly.
+/// # Element-wise q_final parity (not enforced)
 ///
-/// Either way, the contract here is "produce a valid scipy-equivalent
-/// partition", and the existing parity test enforces that.
+/// Switching the parity oracle from partition-equivalence to element-wise
+/// `q_final` would expose this label-permutation gap (qinit columns would
+/// not align). The realistic input distribution and downstream DER are
+/// invariant to relabeling, so this is intentionally not enforced. If a
+/// future test pins element-wise `q_final`, three remediation paths are
+/// available: (1) port scipy's tree-traversal DFS push order verbatim;
+/// (2) compare modulo column permutation recoverable from
+/// `(our_labels, scipy_labels)`; (3) return the permutation alongside
+/// labels and let the caller build a column-permuted qinit.
 fn fcluster_distance_remap(steps: &[Step<f64>], n: usize, threshold: f64) -> Vec<usize> {
   // Single leaf — no merges; one cluster.
   if n == 1 {
@@ -279,18 +274,29 @@ fn fcluster_distance_remap(steps: &[Step<f64>], n: usize, threshold: f64) -> Vec
     }
   }
 
-  // Second pass: scan leaves 0..n and assign encounter-order labels.
-  let mut canonical = vec![0usize; n];
-  let mut next_label = 0usize;
-  let mut label_of_class: HashMap<usize, usize> = HashMap::new();
-  for (i, slot) in canonical.iter_mut().enumerate() {
-    *slot = *label_of_class.entry(raw[i]).or_insert_with(|| {
-      let l = next_label;
-      next_label += 1;
-      l
-    });
-  }
-  canonical
+  // Second pass: `np.unique(raw, return_inverse=True)`-equivalent
+  // canonicalization. Pyannote feeds scipy's `fcluster - 1` through
+  // `np.unique(..., return_inverse=True)` (clustering.py:603-604), which
+  // sorts the distinct DFS-pass labels ascending and remaps each row's
+  // label to its rank in that sorted unique set. The previous
+  // leaf-scan encounter-order canonicalization preserved partition
+  // equivalence but not the label *values*; a downstream caller
+  // (pipeline `assign_embeddings`) builds qinit columns indexed by
+  // these labels, so a value mismatch here produced a column-permuted
+  // qinit, which cascaded into VBx convergence to a different fixed
+  // point on long fixtures (06_long_recording, testaudioset 09/10
+  // and friends). Sorting by raw DFS value matches `np.unique` and
+  // restores bit-exact qinit, q_final, centroid, soft, and
+  // hard_clusters parity downstream.
+  let mut unique_sorted: Vec<usize> = raw.clone();
+  unique_sorted.sort_unstable();
+  unique_sorted.dedup();
+  let value_to_new: HashMap<usize, usize> = unique_sorted
+    .iter()
+    .enumerate()
+    .map(|(i, &v)| (v, i))
+    .collect();
+  raw.iter().map(|v| value_to_new[v]).collect()
 }
 
 /// Recursively assign `label` to every leaf reachable from `node`.

src/cluster/ahc/parity_tests.rs

@@ -192,6 +192,18 @@ fn ahc_init_matches_pyannote_06_long_recording() {
   run_ahc_parity("06_long_recording");
 }
 
+#[test]
+#[ignore = "ad-hoc capture from testaudioset; localizes pyannote parity divergence"]
+fn ahc_init_matches_pyannote_10_mrbeast_clean_water() {
+  run_ahc_parity("10_mrbeast_clean_water");
+}
+
+#[test]
+#[ignore = "ad-hoc capture from testaudioset; localizes 08_luyu_jinjing_freedom +1 spk divergence"]
+fn ahc_init_matches_pyannote_08_luyu_jinjing_freedom() {
+  run_ahc_parity("08_luyu_jinjing_freedom");
+}
+
 /// Remap labels to encounter-order: the first label seen becomes 0,
 /// the second new label becomes 1, etc. After this transform, two
 /// different label arrays representing the same partition compare equal.

src/cluster/ahc/tests.rs

@@ -152,12 +152,23 @@ fn single_row_returns_single_cluster() {
 /// - Row 2 ≈ (0, 1, 0)         → orthogonal
 ///
 /// Distances after L2 norm: d(0,1) ≈ 0.014, d(0,2) ≈ 1.414, d(1,2) ≈ 1.404.
-/// At threshold = 0.5: only the (0,1) pair merges → labels `[0, 0, 1]`.
+/// At threshold = 0.5: only the (0,1) pair merges. Asserts partition
+/// equivalence: rows 0 and 1 share a label, row 2 has a distinct
+/// label. Specific label *values* are determined by
+/// `np.unique`-style canonicalization (sort distinct DFS labels
+/// ascending) and depend on dendrogram traversal.
 #[test]
 fn merges_close_pair_separates_far_row() {
   let m = DMatrix::<f64>::from_row_slice(3, 3, &[1.0, 0.0, 0.0, 100.0, 1.0, 0.0, 0.0, 1.0, 0.0]);
   let labels = ahc_init_dm(&m, 0.5, &crate::ops::spill::SpillOptions::default()).expect("ahc_init");
-  assert_eq!(labels, vec![0, 0, 1]);
+  assert_eq!(
+    labels[0], labels[1],
+    "rows 0 and 1 should share a cluster (got {labels:?})"
+  );
+  assert_ne!(
+    labels[0], labels[2],
+    "row 2 should be its own cluster (got {labels:?})"
+  );
 }
 
 /// All identical rows (after normalization) → single cluster regardless
@@ -183,15 +194,23 @@ fn tiny_threshold_keeps_every_row_isolated() {
   // Three orthogonal directions; pairwise distance after L2 norm ≈ √2 ≈ 1.414.
   let m = DMatrix::<f64>::from_row_slice(3, 3, &[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]);
   let labels = ahc_init_dm(&m, 0.1, &crate::ops::spill::SpillOptions::default()).expect("ahc_init");
-  // Encounter-order labels — each leaf is its own cluster, labelled in
-  // its first-encountered order.
-  assert_eq!(labels, vec![0, 1, 2]);
+  // Each leaf is its own cluster: 3 distinct labels, all from {0, 1, 2}.
+  let mut sorted = labels.clone();
+  sorted.sort_unstable();
+  sorted.dedup();
+  assert_eq!(
+    sorted,
+    vec![0, 1, 2],
+    "expected 3 distinct singleton clusters, got {labels:?}"
+  );
 }
 
-/// Labels must be encounter-order contiguous `0..k` (this is the
-/// `np.unique(return_inverse=True)` post-processing pyannote does).
+/// Labels must be contiguous `0..k` after `np.unique`-style
+/// canonicalization (sort distinct DFS labels ascending). The specific
+/// label values depend on the dendrogram traversal; only partition
+/// equivalence is asserted here.
 #[test]
-fn labels_are_encounter_order_contiguous() {
+fn labels_are_contiguous_after_canonicalization() {
   // Six rows: two pairs that should merge, plus two singletons that
   // shouldn't. Specific arrangement: pair A (rows 0, 3), pair B (rows
   // 1, 4), singleton (row 2), singleton (row 5).
@@ -208,11 +227,23 @@ fn labels_are_encounter_order_contiguous() {
     ],
   );
   let labels = ahc_init_dm(&m, 0.1, &crate::ops::spill::SpillOptions::default()).expect("ahc_init");
-  // Encounter order of labels: row 0 → 0, row 1 → 1, row 2 → 2,
-  // row 3 → 0 (same cluster as row 0), row 4 → 1, row 5 → 3.
-  assert_eq!(labels, vec![0, 1, 2, 0, 1, 3]);
-
-  // Sanity: labels are contiguous 0..k where k = number of distinct.
+  // Partition equivalence: rows 0 and 3 share a cluster, rows 1 and 4
+  // share, rows 2 and 5 are their own clusters.
+  assert_eq!(
+    labels[0], labels[3],
+    "rows 0,3 should share (got {labels:?})"
+  );
+  assert_eq!(
+    labels[1], labels[4],
+    "rows 1,4 should share (got {labels:?})"
+  );
+  assert_ne!(labels[0], labels[1]);
+  assert_ne!(labels[0], labels[2]);
+  assert_ne!(labels[0], labels[5]);
+  assert_ne!(labels[1], labels[2]);
+  assert_ne!(labels[1], labels[5]);
+  assert_ne!(labels[2], labels[5]);
+  // Labels are contiguous 0..k.
   let max = *labels.iter().max().unwrap();
   let mut seen = vec![false; max + 1];
   for &l in &labels {
@@ -277,11 +308,14 @@ fn centroid_linkage_inversion_matches_scipy() {
   //   step 1 (merge 2, {0,1}): d=0.574 ≤ 0.6 BUT subtree's max = 0.65 > 0.6
   //   step 2 (merge 3, ...): d=1.89 > 0.6
   // → no merges accepted; each leaf is its own cluster.
-  // Encounter-order labels: [0, 1, 2, 3].
+  // Each of the 4 leaves is its own cluster: 4 distinct labels.
+  let mut sorted = labels.clone();
+  sorted.sort_unstable();
+  sorted.dedup();
   assert_eq!(
-    labels,
+    sorted,
     vec![0, 1, 2, 3],
-    "inversion case must match scipy: subtree max > threshold means split"
+    "inversion case must match scipy: subtree max > threshold means split (got {labels:?})"
   );
 }