Skip to content

Commit 3c97dc5

Browse files
igerberclaude
andcommitted
fix(rust): deterministic clustered vcov — first-appearance cluster aggregation
The cluster-score aggregation in compute_robust_vcov (also reached via solve_ols(return_vcov=True)) built its (G,k) cluster-scores matrix in HashMap iteration order, which is SipHash-randomized per call: mathematically identical, but the GEMM accumulation order changed on every invocation, wobbling the clustered vcov at ~1e-14 (3 distinct values across 8 identical calls; the Python backend is bit-stable). Rows now accumulate in first-appearance order — ascending for the factorized 0..G-1 ids the Python dispatcher passes, matching NumPy's groupby order. Verified 1 distinct value across 50 identical calls; NumPy parity unchanged (~1e-14 cross-backend GEMM tolerance). Bit-identity regression tests cover contiguous and non-contiguous unsorted cluster ids. Resolves the TODO row currently in flight on PR #647 (row deletion follows once both merge). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent c4aaab2 commit 3c97dc5

3 files changed

Lines changed: 62 additions & 12 deletions

File tree

CHANGELOG.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
3535
`did_had_pretest_workflow`, ...) are unchanged in this release and removed separately.
3636

3737
### Fixed
38+
- **Rust clustered vcov is now run-to-run deterministic.** The cluster-score aggregation
39+
in `compute_robust_vcov` (also reached via `solve_ols(return_vcov=True)`) built its
40+
(G, k) cluster-scores matrix in `HashMap` iteration order, which is SipHash-randomized
41+
per call — mathematically identical, but the GEMM accumulation order changed on every
42+
invocation, wobbling the vcov at ~1e-14 (3 distinct values observed across 8 identical
43+
calls; the Python backend was bit-stable). Rows now accumulate in first-appearance
44+
order — ascending for the factorized 0..G-1 ids the Python dispatcher passes, matching
45+
NumPy's groupby order. Verified: 1 distinct value across 50 identical calls (was 3/8);
46+
NumPy parity unchanged (~1e-14, the normal cross-backend GEMM tolerance); bit-identity
47+
regression tests cover both contiguous and non-contiguous unsorted cluster ids.
3848
- **`HonestDiD` Δ^SD optimal-FLCI center parity with R (SE-audit B2b).** The optimal
3949
Fixed-Length CI optimizer was a flat Nelder-Mead over slope weights that landed on a
4050
different affine estimator than R `HonestDiD::findOptimalFLCI` at intermediate smoothness

rust/src/linalg.rs

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -349,18 +349,22 @@ fn compute_robust_vcov_internal(
349349
let residuals_col = residuals.insert_axis(Axis(1)); // (n, 1)
350350
let scores = x * &residuals_col; // (n, k) - broadcasts residuals across columns
351351

352-
// Aggregate scores by cluster using HashMap
353-
let mut cluster_sums: HashMap<i64, Array1<f64>> = HashMap::new();
352+
// Aggregate scores by cluster DETERMINISTICALLY. HashMap
353+
// iteration order is SipHash-randomized per map instance, which
354+
// reordered the cluster rows on every call — mathematically
355+
// identical, but the GEMM accumulation order changed, making
356+
// the clustered vcov run-to-run nondeterministic at ~1e-14
357+
// (distinct values across identical calls; the Python backend
358+
// is bit-stable). Rows accumulate in first-appearance order
359+
// instead: for the factorized 0..G-1 ids the Python dispatcher
360+
// passes this is ascending id order, matching NumPy's groupby.
361+
let mut cluster_index: HashMap<i64, usize> = HashMap::new();
354362
for i in 0..n_obs {
355-
let cluster = clusters[i];
356-
let row = scores.row(i).to_owned();
357-
cluster_sums
358-
.entry(cluster)
359-
.and_modify(|sum| *sum = &*sum + &row)
360-
.or_insert(row);
363+
let next = cluster_index.len();
364+
cluster_index.entry(clusters[i]).or_insert(next);
361365
}
362366

363-
let n_clusters = cluster_sums.len();
367+
let n_clusters = cluster_index.len();
364368

365369
if n_clusters < 2 {
366370
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
@@ -369,10 +373,13 @@ fn compute_robust_vcov_internal(
369373
)));
370374
}
371375

372-
// Build cluster scores matrix (G, k)
376+
// Build cluster scores matrix (G, k) in first-appearance order
373377
let mut cluster_scores = Array2::<f64>::zeros((n_clusters, k));
374-
for (idx, (_cluster_id, sum)) in cluster_sums.iter().enumerate() {
375-
cluster_scores.row_mut(idx).assign(sum);
378+
for i in 0..n_obs {
379+
let idx = cluster_index[&clusters[i]];
380+
cluster_scores
381+
.row_mut(idx)
382+
.zip_mut_with(&scores.row(i), |a, b| *a += *b);
376383
}
377384

378385
// Compute meat: Σ_g (X_g' e_g)(X_g' e_g)'

tests/test_rust_backend.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3622,3 +3622,36 @@ def _unstable(*a, **k):
36223622
v = la.compute_robust_vcov(X, resid, vcov_type="hc2")
36233623
v_py = la._compute_robust_vcov_numpy(X, resid, None, vcov_type="hc2")
36243624
np.testing.assert_allclose(v, v_py, rtol=1e-12, atol=1e-15)
3625+
class TestClusterVcovDeterminism:
3626+
"""Clustered vcov is bit-identical across repeated identical calls.
3627+
3628+
The cluster-score aggregation previously built rows in HashMap iteration
3629+
order (SipHash-randomized per call): mathematically identical, but the
3630+
GEMM accumulation order changed run-to-run, wobbling the vcov at ~1e-14
3631+
(3 distinct values observed in 8 identical calls) while the Python
3632+
backend was bit-stable. Rows now accumulate in first-appearance order
3633+
(ascending for the factorized ids the dispatcher passes)."""
3634+
3635+
def test_solve_ols_cluster_vcov_bit_identical_across_calls(self):
3636+
from diff_diff.linalg import solve_ols
3637+
3638+
rng = np.random.default_rng(0)
3639+
X = rng.normal(size=(400, 3))
3640+
y = rng.normal(size=400)
3641+
cl = np.repeat(np.arange(10), 40)
3642+
baseline = solve_ols(X, y, cluster_ids=cl, return_vcov=True)[2]
3643+
for _ in range(20):
3644+
v = solve_ols(X, y, cluster_ids=cl, return_vcov=True)[2]
3645+
np.testing.assert_array_equal(v, baseline)
3646+
3647+
def test_compute_robust_vcov_cluster_bit_identical_across_calls(self):
3648+
from diff_diff.linalg import compute_robust_vcov
3649+
3650+
rng = np.random.default_rng(1)
3651+
X = rng.normal(size=(300, 4))
3652+
resid = rng.normal(size=300)
3653+
# Non-contiguous, unsorted ids exercise the first-appearance remap.
3654+
cl = np.repeat(np.array([7, 3, 11, 5, 42, 3, 7, 11, 5, 42]), 30)
3655+
baseline = compute_robust_vcov(X, resid, cluster_ids=cl)
3656+
for _ in range(20):
3657+
np.testing.assert_array_equal(compute_robust_vcov(X, resid, cluster_ids=cl), baseline)

0 commit comments

Comments
 (0)