Repository navigation
[MRG] Make BSP-OT random seeding reproducible across calls - #898
xihaian251 wants to merge 3 commits into
Conversation
ot.bsp.compute_bspot_bijection documents `seed` as a reproducibility knob, but the BSP-OT engines are `thread_local std::mt19937` objects constructed once per thread from `global_seed` (BSP-OT_header_only.h:549-553), so the `global_seed = seed` assignment in BSPOT_wrap reaches the engine of the first call made on that thread only: a repeated call continued the previous stream and returned different plans, different intermediary stacks and a different cost (Issue PythonOT#897). The same lines also cost solution quality, because the candidate plans are generated by `#pragma omp parallel for` and each worker built its engine from the same `global_seed`: iterations at equal positions in their chunks consumed identical subsequences and produced identical matchings, so `n_plans=64` yielded ceil(n_plans / n_threads) distinct candidates on this machine and `MergePlansNoPar` walked over the copies. Re-seed the engine at the start of each iteration from a SplitMix64 mix of (user_seed, plan index) and pass the seed down as a parameter, so that candidate plan i is a function of the arguments alone. The mixing must not be additive: with `user_seed + task`, plan t of seed s is plan t-1 of seed s+1, so changing the seed would leave nearly the whole candidate set unchanged. The matching object is now constructed per iteration, which removes the only mutable state shared between the iterations of this loop, and the now dead `global_seed = seed` write is dropped; those two thread-safety edits do not change the result for a given seed.
Repeated calls, interleaved history, a fresh-interpreter reference, invariance to OMP_NUM_THREADS, and the property that candidate plan i is determined by (seed, i), which is checked by comparing the first rows of the candidate stacks returned for n_plans=4 and n_plans=37. Five of the six fail on current master. The candidate matchings are deliberately not asserted to be pairwise distinct: independent streams can legitimately return the same bijection -- for a similarity transform every BSP candidate is the identity -- so distinctness is a property of the instance and of the standard library's distributions rather than of the seeding, and PythonOT#897 reports it as a quality metric. No historical cost or permutation is hardcoded, so the tests do not freeze the new random sequence. The children report the extension hash they loaded and the parent asserts it matches, so a child importing another build cannot silently pass.
Required by .github/workflows/check_release.yml, which fails any pull request that does not touch RELEASES.md.
|
The red checks are an infrastructure failure, not a test failure: every failed or cancelled job died in No job failed at an install, build, lint or test step. For reference, the two jobs that did get past checkout are green on head
with |
|
Hi maintainers, Could someone with write access please re-run the failed jobs for PR #898? The failures occurred during Windows and macOS have already passed. Affected runs:
A re-run would help verify whether there are any actual code or test failures. Thanks for your help! |
Why
ot.bsp.compute_bspot_bijection(..., seed=...)— and thereforerandom_stateofot.solve_sample(..., method="bsp")— is documented as a reproducibility knob, but it only held for thefirst call made on a thread. Reported with measurements in #897.
engis a thread-local object whose constructor consumesglobal_seedthe first timerng()is used onthat thread and never again, so on later calls
global_seed = seedcannot reach it and every drawcontinues the previous stream. Three identical calls in one process returned three different bijections
and three different costs (0.786970722 / 0.786334761 / 0.793428761 on the instance in #897).
The same lines also cost solution quality: the candidate matchings are produced by
#pragma omp parallel for(ot/bsp/bsp_wrapper.cpp:40-48) and each worker builds its engine from thesame
global_seed, so with the default static schedule the iterations at equal positions in theirchunks consume identical subsequences and return identical matchings. Measured,
n_plans=64:64 distinct candidates at 1 thread, 32 at 2, 16 at 4, 8 at 8, 4 at 16 — i.e.
ceil(n_plans / n_threads)distinct plans, the rest being copies that
MergePlansNoParthen walks over again.What
Both changes are inside the loop that generates the candidate plans, in
ot/bsp/bsp_wrapper.cpp:thread_local) engine at the start of each iteration with a value derived from(user_seed, plan index), and passseeddown as a parameter instead of reading the mutable global;BijectiveBSPMatchingper iteration, so the loop body no longer writes state sharedwith other iterations (
random_pivotwas written by every thread of the team).The derivation must not be additive: with
user_seed + task, plantof seedsis plant-1of seeds+1, so changing the seed would leave nearly the whole candidate set unchanged — hence the SplitMix64finalizer. The vendored amalgamated header, the Cython layer, the Python API and every default are
untouched.
Contract
For a fixed build,
compute_bspot_bijection(X, Y, n_plans, p, initial_perm, gaussian_slicing, seed)isnow a function of its arguments: repeat-equality, equality after unrelated calls, equality with a first
call in a fresh interpreter, and equality across
OMP_NUM_THREADS. Candidate planidepends on(seed, i)alone, which is what the new "seed and index" test checks by comparing the first rows of thestacks returned for
n_plans=4andn_plans=37.Two deliberate non-claims. Results are reproducible per build, not across standard libraries:
std::mt19937::seedis specified by the standard, butstd::normal_distributionandstd::uniform_int_distributionare implementation-specified, so a guarantee across libstdc++/libc++/MSVCwould need self-contained distributions. And different seeds are not required to give different
bijections. For the same reason the tests do not assert that the candidate matchings are pairwise
distinct — independent streams can legitimately return the same bijection (for
Y = 2X + tevery BSPcandidate is the identity, and on a 20-point instance I measured 120 different seeds producing only 114
distinct matchings). Diversity of the candidates is the consequence that #897 reports as a quality
metric (4/64 → 64/64 here), not something a test should promise.
For a given seed the returned permutation does change; unavoidable, since the old sequence depended on
the call history and on the CPU count, so there was no stable behaviour to keep. Existing assertions on
fixed configurations (
test_bsp_ot_exact_identity,test_bsp_ot_identity_null_cost,test_bsp_ot_bijective,test_bsp_ot_torch_backend) hold for any stream and still pass.Added cost: one
mt19937state initialisation (624 words) per candidate plan, i.e.O(1)innagainst the
O(n log n)work of generating that plan.Tests
Six tests added to
test/test_bspot.py; no historical cost or permutation is hardcoded, so they do notfreeze the new random sequence. Two of them spawn child interpreters (cross-process reproducibility and
invariance to
OMP_NUM_THREADSis otherwise untestable); the child also reports the SHA-256 of theextension it loaded, and the parent asserts it matches its own, so a child importing a different build
cannot silently pass.
Measured on Windows 11, MSVC 14.44 with
/openmp, 16 logical CPUs, Python 3.13.1, numpy 2.4.6,torch 2.10.0+cpu, jax 0.12.0, in a fresh clone at
5cb82c9, built withpython setup.py build_ext --inplace:correct — it is not the defect);
OMP_NUM_THREADS=1, 2, 4, 8, 16;pytest test/test_solvers.py -k bsp: 6 passed;n_plans1–256, bothslicing modes, several seeds; each also compared against the exact
ot.emdvalue) checksrepeat-equality, history-equality, permutation structure, and consistency between the cost the solver
reports for its own permutation and the cost of that permutation: 39 configurations pass, 0 fail.
The same script against the master build fails every one of the 39.
ruff checkon the changed test file reports the same 5 pre-existing findings asHEAD, so no newlint.
Known limitation, stated rather than hidden: on a build or machine where OpenMP is unavailable (POT's
setup.pydrops the flag silently) or pinned to one thread, the thread-count test cannot distinguishthe old behaviour from the new one; and a change that simply deleted
#pragma omp parallel forwouldalso pass these tests — staying parallel is a performance property, not something these assertions
cover.
Out of scope
n_plans=0crashesMergePlansNoPar(it indexes an empty plan vector), andmerge_bijectionsreadsplans.shape[1]as the number of plans, which is an out-of-bounds read whenn_plans < n. Both areseparate bugs; I will report them separately instead of mixing them into this change.
Fixes #897