From 07b7ff7656a7ad56911b42d9fd535e0280e4127d Mon Sep 17 00:00:00 2001 From: xihaian251 Date: Tue, 3 Nov 2026 13:01:01 +0800 Subject: [PATCH 1/5] Read the number of plans from axis 0 in merge_bijections perms is a stack of (n_plans, n_samples) permutations, but merge_bijections_c passed plans.shape[1] -- the number of samples -- as the plan count, so the C++ merge loop walked n_samples plans over a buffer holding n_plans. With n_plans < n_samples it read past the end of the array (heap-buffer-overflow under AddressSanitizer, segmentation fault through the public API); with n_plans > n_samples it silently dropped every plan past the n_samples-th. Take the count from plans.shape[0] and reject a perms that does not follow the documented layout at the Python boundary, so a wrong row length raises instead of reaching C++. --- ot/bsp/bsp_ot.py | 11 +++++ ot/bsp/bsp_wrap.pyx | 5 +- test/test_bspot.py | 113 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 128 insertions(+), 1 deletion(-) diff --git a/ot/bsp/bsp_ot.py b/ot/bsp/bsp_ot.py index d7d13a7b5..6c32793d7 100644 --- a/ot/bsp/bsp_ot.py +++ b/ot/bsp/bsp_ot.py @@ -152,6 +152,17 @@ def merge_bijections(X, Y, perms, p=2): perms = np.asarray(perms, dtype=np.int32) + if perms.ndim != 2: + raise ValueError("perms must be a 2D array of shape (n_plans, n_samples)") + + if perms.shape[1] != X_np.shape[0]: + raise ValueError( + "perms must have shape (n_plans, n_samples) where n_samples is the number " + "of points in X and Y, got {} plans of length {} for {} points".format( + perms.shape[0], perms.shape[1], X_np.shape[0] + ) + ) + cost_inner, perm = merge_bijections_c(X_np, Y_np, perms, p) perm = nx.from_numpy(perm) diff --git a/ot/bsp/bsp_wrap.pyx b/ot/bsp/bsp_wrap.pyx index 043028832..2a3659cbb 100644 --- a/ot/bsp/bsp_wrap.pyx +++ b/ot/bsp/bsp_wrap.pyx @@ -81,7 +81,10 @@ def merge_bijections_c(np.ndarray[double, ndim=2, mode="c"] X, np.ndarray[double cdef int n = X.shape[0] cdef int d = X.shape[1] - cdef int k = plans.shape[1] + # plans is laid out as (n_plans, n_samples), so the number of plans to merge is + # shape[0]. Reading shape[1] instead makes the C++ loop walk n_samples plans over a + # buffer that only holds n_plans, which reads past it when n_plans < n_samples. + cdef int k = plans.shape[0] cdef np.ndarray[int, ndim=1, mode="c"] plan = np.zeros(n, dtype=np.int32) cdef double cost_val diff --git a/test/test_bspot.py b/test/test_bspot.py index 862923e12..fe058a50a 100644 --- a/test/test_bspot.py +++ b/test/test_bspot.py @@ -4,6 +4,8 @@ # # License: MIT License +import warnings + import ot import ot.bsp import numpy as np @@ -135,3 +137,114 @@ def test_bsp_ot_torch_backend(): # cost should be zero np.testing.assert_allclose(cost_new, 0, atol=1e-5) + + +def _mean_squared_cost(A, B, perm): + # cost of ``perm`` computed here, independently of what the solver reports + return float(np.mean(np.sum((A - B[np.asarray(perm)]) ** 2, axis=1))) + + +def test_bsp_merge_fewer_plans_than_samples(): + # perms is a stack of (n_plans, n_samples) permutations, so the merge has n_plans of + # them to read. With one plan over 32 samples the old code read one plan per sample + # and walked off the end of the array (issue #910). + n = 32 + rng = np.random.RandomState(0) + + A = rng.randn(n, 2) + B = rng.randn(n, 2) + perms = np.arange(n, dtype=np.int32).reshape(1, n) + + with warnings.catch_warnings(): + # merge_bijections warns when its reported cost does not match the permutation + warnings.simplefilter("error") + cost, perm = ot.bsp.merge_bijections(A, B, perms, 2) + + perm = np.asarray(perm) + np.testing.assert_array_equal(np.sort(perm), np.arange(n)) + np.testing.assert_allclose(cost, _mean_squared_cost(A, B, perm), rtol=1e-9) + assert cost <= _mean_squared_cost(A, B, np.arange(n)) + 1e-12 + + +def test_bsp_merge_uses_every_supplied_plan(): + # The plans past the n_samples-th one must reach the merge. Both clouds are the same, + # so the identity is the only plan with zero cost and the only one that can produce + # it: if it is dropped the merged cost stays on the mediocre plan. + n = 8 + rng = np.random.RandomState(0) + + A = rng.randn(n, 2) + B = A.copy() + + mediocre = np.roll(np.arange(n), 1).astype(np.int32) + without = np.tile(mediocre, (n, 1)) + with_last = np.vstack([without, np.arange(n)]).astype(np.int32) + + cost_without, _ = ot.bsp.merge_bijections(A, B, without, 2) + cost_with, perm_with = ot.bsp.merge_bijections(A, B, with_last, 2) + + assert cost_with < cost_without, "the plan past n_samples was ignored" + np.testing.assert_array_equal(np.asarray(perm_with), np.arange(n)) + np.testing.assert_allclose(cost_with, 0.0, atol=1e-12) + + +def test_bsp_merge_rejects_plans_of_wrong_length(): + # A stack whose rows are not one entry per point cannot be merged; it must fail here + # instead of being read as a number of plans by C++. + n = 8 + rng = np.random.RandomState(0) + + A = rng.randn(n, 2) + B = rng.randn(n, 2) + + with pytest.raises(ValueError, match="n_samples is the number of points"): + ot.bsp.merge_bijections(A, B, np.tile(np.roll(np.arange(7), 1), (3, 1)), 2) + + with pytest.raises(ValueError, match="2D"): + ot.bsp.merge_bijections(A, B, np.arange(n, dtype=np.int32), 2) + + with pytest.raises(ValueError, match="2D"): + ot.bsp.merge_bijections(A, B, np.zeros((n, n, 1), dtype=np.int32), 2) + + +def test_bsp_merge_on_a_standard_stack(): + # The ordinary case must keep its documented behaviour: a bijection no worse than any + # plan it merged, and a cost that is the cost of that bijection. + n = 50 + rng = np.random.RandomState(0) + + A = rng.randn(n, 2) + B = rng.randn(n, 2) + perms = np.stack([np.roll(np.arange(n), s) for s in (1, 7, 23, 41)]).astype(np.int32) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + cost, perm = ot.bsp.merge_bijections(A, B, perms, 2) + + perm = np.asarray(perm) + np.testing.assert_array_equal(np.sort(perm), np.arange(n)) + np.testing.assert_allclose(cost, _mean_squared_cost(A, B, perm), rtol=1e-9) + for plan in perms: + assert cost <= _mean_squared_cost(A, B, plan) + 1e-12 + + +def test_bsp_merge_converts_dtypes_and_refuses_views(): + # np.asarray(perms, dtype=np.int32) already accepted these, and the fix must not change + # that; a non-contiguous view was refused before the fix too, rather than read loosely. + n = 8 + rng = np.random.RandomState(0) + + A = rng.randn(n, 2) + B = rng.randn(n, 2) + plans = np.tile(np.roll(np.arange(n), 1), (4, 1)) + + ref_cost, ref_perm = ot.bsp.merge_bijections(A, B, plans.astype(np.int32), 2) + for other in (plans.astype(np.int64), plans.astype(np.float64), plans.tolist()): + cost, perm = ot.bsp.merge_bijections(A, B, other, 2) + np.testing.assert_allclose(cost, ref_cost, rtol=0, atol=0) + np.testing.assert_array_equal(np.asarray(perm), np.asarray(ref_perm)) + + view = np.asfortranarray(plans.astype(np.int32)) + assert not view.flags["C_CONTIGUOUS"] + with pytest.raises(ValueError, match="C-contiguous"): + ot.bsp.merge_bijections(A, B, view, 2) From 968ff2e38b1f6782f1a0dd442c36e98d56fa4cdd Mon Sep 17 00:00:00 2001 From: xihaian251 Date: Tue, 3 Nov 2026 13:41:10 +0800 Subject: [PATCH 2/5] Refuse a malformed plan stack at the native boundary too AddressSanitizer on the built extension still reports a heap-buffer-overflow when merge_bijections_c is called directly with rows that are not one entry per point: the C++ loop reads n_plans * n_samples integers, so a shorter row leaves it short of the buffer even with the plan count read from the right axis. Check plans.shape[1] against the number of points in the Cython layer and assert that refusal from the test as well, so the overrun is unreachable for any caller rather than only for the public wrapper. --- ot/bsp/bsp_wrap.pyx | 7 +++++++ test/test_bspot.py | 13 +++++++++++++ 2 files changed, 20 insertions(+) diff --git a/ot/bsp/bsp_wrap.pyx b/ot/bsp/bsp_wrap.pyx index 2a3659cbb..7093cffab 100644 --- a/ot/bsp/bsp_wrap.pyx +++ b/ot/bsp/bsp_wrap.pyx @@ -85,6 +85,13 @@ def merge_bijections_c(np.ndarray[double, ndim=2, mode="c"] X, np.ndarray[double # shape[0]. Reading shape[1] instead makes the C++ loop walk n_samples plans over a # buffer that only holds n_plans, which reads past it when n_plans < n_samples. cdef int k = plans.shape[0] + + # C++ reads k * n integers, so a row that is shorter or longer than the point cloud + # would still make it run past the buffer even with the right plan count. + if plans.shape[1] != n: + raise ValueError( + "plans must have shape (n_plans, n_samples) with one entry per point") + cdef np.ndarray[int, ndim=1, mode="c"] plan = np.zeros(n, dtype=np.int32) cdef double cost_val diff --git a/test/test_bspot.py b/test/test_bspot.py index fe058a50a..eb8a70cc4 100644 --- a/test/test_bspot.py +++ b/test/test_bspot.py @@ -206,6 +206,19 @@ def test_bsp_merge_rejects_plans_of_wrong_length(): with pytest.raises(ValueError, match="2D"): ot.bsp.merge_bijections(A, B, np.zeros((n, n, 1), dtype=np.int32), 2) + # the same refusal at the native boundary, so no caller can reach the overrun + from ot.bsp import bsp_wrap + + with pytest.raises(ValueError, match="one entry per point"): + bsp_wrap.merge_bijections_c( + np.ascontiguousarray(A, dtype=np.float64), + np.ascontiguousarray(B, dtype=np.float64), + np.ascontiguousarray( + np.tile(np.roll(np.arange(7), 1), (3, 1)), dtype=np.int32 + ), + 2, + ) + def test_bsp_merge_on_a_standard_stack(): # The ordinary case must keep its documented behaviour: a bijection no worse than any From 615b28c5f2d48c13fd48035c4d394210b5a167fc Mon Sep 17 00:00:00 2001 From: xihaian251 Date: Tue, 3 Nov 2026 13:58:50 +0800 Subject: [PATCH 3/5] Note the merge_bijections plan-count fix in RELEASES.md --- RELEASES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/RELEASES.md b/RELEASES.md index 5d58d2128..f9bb3d8f9 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -18,6 +18,7 @@ #### Closed issues +- Fix an out-of-bounds read in `ot.bsp.merge_bijections`: the number of plans to merge was taken from `perms.shape[1]`, the number of samples, instead of `perms.shape[0]`, so the C++ merge loop walked `n_samples` plans over a stack of `n_plans` and read past it when `n_plans < n_samples`, or dropped the plans at index `>= n_samples` otherwise. A `perms` that does not follow the documented `(n_plans, n_samples)` layout is now rejected before it reaches C++ (PR #911, Issue #910) - Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264) - Fix duplicate function and class descriptions in generated module API pages while preserving top-level aliases and example backreferences (Issue #527) - Fix internal arrays created on the torch default device instead of the input device in `ot.partial.partial_wasserstein`, `ot.wasserstein_circle` (and `ot.sliced_wasserstein_sphere`), the high-dimensional Bures-Wasserstein functions in `ot.gaussian`, `ot.gmm`, `ot.lowrank_sinkhorn` with `init="deterministic"`, the semi-relaxed (F)GW barycenters and `ot.utils.projection_sparse_simplex`, which made them fail on GPU inputs. Add a `meta_default_device` test fixture to catch such allocations on CPU (PR #883, Issue #852) From b15c1572735cdf41831d6b68f1956b050db9815a Mon Sep 17 00:00:00 2001 From: B Date: Tue, 3 Nov 2026 19:38:13 +0800 Subject: [PATCH 4/5] fix: validate BSP bijection plans before merging merge_bijections accepted any stack whose rows were one entry per point, so a row that is not a permutation of 0..n_samples-1 reached BijectiveMatching and indexed the point clouds out of bounds, and an empty stack left the merge helper with nothing to read. Both crash the interpreter. Reject an empty plan stack and check that every row is a permutation of 0..n_samples-1, at the Python entry point and again in merge_bijections_c so a direct caller cannot reach C++ with a malformed stack. The entries are checked before the conversion to int32, which truncates fractional values and wraps integers that do not fit and would otherwise turn an invalid row into a valid one. --- RELEASES.md | 2 +- ot/bsp/bsp_ot.py | 15 ++++++- ot/bsp/bsp_wrap.pyx | 26 ++++++++++--- test/test_bspot.py | 95 +++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 131 insertions(+), 7 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index f9bb3d8f9..e4234e508 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -18,7 +18,7 @@ #### Closed issues -- Fix an out-of-bounds read in `ot.bsp.merge_bijections`: the number of plans to merge was taken from `perms.shape[1]`, the number of samples, instead of `perms.shape[0]`, so the C++ merge loop walked `n_samples` plans over a stack of `n_plans` and read past it when `n_plans < n_samples`, or dropped the plans at index `>= n_samples` otherwise. A `perms` that does not follow the documented `(n_plans, n_samples)` layout is now rejected before it reaches C++ (PR #911, Issue #910) +- Fix an out-of-bounds read in `ot.bsp.merge_bijections`: the number of plans to merge was taken from `perms.shape[1]`, the number of samples, instead of `perms.shape[0]`, so the C++ merge loop walked `n_samples` plans over a stack of `n_plans` and read past it when `n_plans < n_samples`, or dropped the plans at index `>= n_samples` otherwise. A `perms` that does not follow the documented `(n_plans, n_samples)` layout is now rejected before it reaches C++, as are an empty plan stack and rows that are not a permutation of `0..n_samples-1`, at the Python and at the native entry point (PR #911, Issue #910) - Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264) - Fix duplicate function and class descriptions in generated module API pages while preserving top-level aliases and example backreferences (Issue #527) - Fix internal arrays created on the torch default device instead of the input device in `ot.partial.partial_wasserstein`, `ot.wasserstein_circle` (and `ot.sliced_wasserstein_sphere`), the high-dimensional Bures-Wasserstein functions in `ot.gaussian`, `ot.gmm`, `ot.lowrank_sinkhorn` with `init="deterministic"`, the semi-relaxed (F)GW barycenters and `ot.utils.projection_sparse_simplex`, which made them fail on GPU inputs. Add a `meta_default_device` test fixture to catch such allocations on CPU (PR #883, Issue #852) diff --git a/ot/bsp/bsp_ot.py b/ot/bsp/bsp_ot.py index 6c32793d7..caefe8ea0 100644 --- a/ot/bsp/bsp_ot.py +++ b/ot/bsp/bsp_ot.py @@ -150,7 +150,7 @@ def merge_bijections(X, Y, perms, p=2): if X_np.ndim != 2: raise ValueError("X and Y must be 2D arrays") - perms = np.asarray(perms, dtype=np.int32) + perms = np.asarray(perms) if perms.ndim != 2: raise ValueError("perms must be a 2D array of shape (n_plans, n_samples)") @@ -163,6 +163,19 @@ def merge_bijections(X, Y, perms, p=2): ) ) + if perms.shape[0] == 0: + raise ValueError("perms must contain at least one plan") + + # Check the entries before casting to int32: the cast truncates fractional values and + # wraps integers that do not fit, either of which can turn an invalid row into a valid + # one. The merge indexes the point clouds with these values, so they must be a + # permutation of 0..n_samples-1 first. + expected = np.broadcast_to(np.arange(X_np.shape[0]), perms.shape) + if not np.array_equal(np.sort(perms, axis=1), expected): + raise ValueError("each row of perms must be a permutation of 0..n_samples-1") + + perms = np.asarray(perms, dtype=np.int32) + cost_inner, perm = merge_bijections_c(X_np, Y_np, perms, p) perm = nx.from_numpy(perm) diff --git a/ot/bsp/bsp_wrap.pyx b/ot/bsp/bsp_wrap.pyx index 7093cffab..38c2e6448 100644 --- a/ot/bsp/bsp_wrap.pyx +++ b/ot/bsp/bsp_wrap.pyx @@ -81,17 +81,33 @@ def merge_bijections_c(np.ndarray[double, ndim=2, mode="c"] X, np.ndarray[double cdef int n = X.shape[0] cdef int d = X.shape[1] - # plans is laid out as (n_plans, n_samples), so the number of plans to merge is - # shape[0]. Reading shape[1] instead makes the C++ loop walk n_samples plans over a - # buffer that only holds n_plans, which reads past it when n_plans < n_samples. cdef int k = plans.shape[0] - # C++ reads k * n integers, so a row that is shorter or longer than the point cloud - # would still make it run past the buffer even with the right plan count. if plans.shape[1] != n: raise ValueError( "plans must have shape (n_plans, n_samples) with one entry per point") + if k == 0: + raise ValueError("plans must contain at least one plan") + + # C++ indexes the point clouds with these entries, so a stack whose rows are not a + # permutation of 0..n-1 must not reach BijectiveMatching. Counting is linear in the + # size of the stack and does not copy it. + cdef np.ndarray[int, ndim=1, mode="c"] seen = np.zeros(n, dtype=np.int32) + cdef int i + cdef int j + cdef int v + + for i in range(k): + for j in range(n): + seen[j] = 0 + for j in range(n): + v = plans[i, j] + if v < 0 or v >= n or seen[v] != 0: + raise ValueError( + "each row of plans must be a permutation of 0..n_samples-1") + seen[v] = 1 + cdef np.ndarray[int, ndim=1, mode="c"] plan = np.zeros(n, dtype=np.int32) cdef double cost_val diff --git a/test/test_bspot.py b/test/test_bspot.py index eb8a70cc4..df5b857f4 100644 --- a/test/test_bspot.py +++ b/test/test_bspot.py @@ -261,3 +261,98 @@ def test_bsp_merge_converts_dtypes_and_refuses_views(): assert not view.flags["C_CONTIGUOUS"] with pytest.raises(ValueError, match="C-contiguous"): ot.bsp.merge_bijections(A, B, view, 2) + + +def test_bsp_merge_rejects_empty_plan_stack(): + # An empty stack leaves the merge helper with nothing to read, which crashes the + # interpreter, so it has to be refused before the helper is reached. + n = 8 + rng = np.random.RandomState(0) + + A = rng.randn(n, 2) + B = rng.randn(n, 2) + + with pytest.raises(ValueError, match="at least one plan"): + ot.bsp.merge_bijections(A, B, np.zeros((0, n), dtype=np.int32), 2) + + # the native entry point is reachable on its own, so it must refuse it too + from ot.bsp import bsp_wrap + + with pytest.raises(ValueError, match="at least one plan"): + bsp_wrap.merge_bijections_c( + np.ascontiguousarray(A), + np.ascontiguousarray(B), + np.zeros((0, n), dtype=np.int32), + 2, + ) + + +def test_bsp_merge_rejects_rows_that_are_not_permutations(): + # A row describes a bijection only if it is a permutation of 0..n-1. Anything else + # makes the merge index outside the point clouds, so it is refused rather than read. + n = 8 + rng = np.random.RandomState(0) + + A = rng.randn(n, 2) + B = rng.randn(n, 2) + identity = np.arange(n, dtype=np.int32) + + invalid = [ + np.array([0, 0, 2, 3, 4, 5, 6, 7], dtype=np.int32), # repeated entry + np.array([-1, 1, 2, 3, 4, 5, 6, 7], dtype=np.int32), # negative index + np.array([0, 1, 2, 3, 4, 5, 6, 8], dtype=np.int32), # past the last point + np.array([0, 1, 2, 3, 4, 5, 6, 6], dtype=np.int32), # missing entry + ] + + from ot.bsp import bsp_wrap + + for row in invalid: + stack = row.reshape(1, n) + with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): + ot.bsp.merge_bijections(A, B, stack, 2) + with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): + bsp_wrap.merge_bijections_c( + np.ascontiguousarray(A), + np.ascontiguousarray(B), + np.ascontiguousarray(stack, dtype=np.int32), + 2, + ) + + # one bad row is enough to refuse the whole stack, wherever it sits in it + mixed = np.stack([identity, row, identity]) + with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): + ot.bsp.merge_bijections(A, B, mixed, 2) + with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): + bsp_wrap.merge_bijections_c( + np.ascontiguousarray(A), + np.ascontiguousarray(B), + np.ascontiguousarray(mixed, dtype=np.int32), + 2, + ) + + +def test_bsp_merge_rejects_plans_that_only_look_valid_after_casting(): + # The entries are checked before the conversion to int32. That cast truncates + # fractional values and wraps integers that do not fit, and either one would silently + # turn an invalid row into a permutation if the check ran after it. + n = 8 + rng = np.random.RandomState(0) + + A = rng.randn(n, 2) + B = rng.randn(n, 2) + + fractional = np.arange(n, dtype=np.float64).reshape(1, n) + fractional[0, 3] = 3.5 + with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): + ot.bsp.merge_bijections(A, B, fractional, 2) + + wrapped = np.arange(n, dtype=np.int64).reshape(1, n) + wrapped[0, 3] = 2**32 + 3 + assert wrapped.astype(np.int32)[0, 3] == 3, "the cast keeps the low bits only" + with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): + ot.bsp.merge_bijections(A, B, wrapped, 2) + + nan = np.arange(n, dtype=np.float64).reshape(1, n) + nan[0, 5] = np.nan + with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): + ot.bsp.merge_bijections(A, B, nan, 2) From fa8f087c0cae2256648bfb7a87166b1d6848f5c7 Mon Sep 17 00:00:00 2001 From: B Date: Wed, 4 Nov 2026 22:07:48 +0800 Subject: [PATCH 5/5] Drop the native permutation loop and trim per review merge_bijections_c keeps only the guards that bound the buffer handed to C++ (row length, empty stack); the Python wrapper validates the permutation values. Invalid-input tests are parameterised, and the RELEASES.md entry is shortened. --- RELEASES.md | 2 +- ot/bsp/bsp_ot.py | 6 +- ot/bsp/bsp_wrap.pyx | 18 ---- test/test_bspot.py | 225 +++++++++++++++----------------------------- 4 files changed, 79 insertions(+), 172 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index c51ddc04b..85be4654d 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -18,7 +18,7 @@ #### Closed issues -- Fix an out-of-bounds read in `ot.bsp.merge_bijections`: the number of plans to merge was taken from `perms.shape[1]`, the number of samples, instead of `perms.shape[0]`, so the C++ merge loop walked `n_samples` plans over a stack of `n_plans` and read past it when `n_plans < n_samples`, or dropped the plans at index `>= n_samples` otherwise. A `perms` that does not follow the documented `(n_plans, n_samples)` layout is now rejected before it reaches C++, as are an empty plan stack and rows that are not a permutation of `0..n_samples-1`, at the Python and at the native entry point (PR #911, Issue #910) +- Fix an out-of-bounds read in `ot.bsp.merge_bijections`: the number of plans was taken from `perms.shape[1]` (the number of samples) instead of `perms.shape[0]`, so merging fewer plans than samples read past the array and merging more silently dropped plans. A `perms` outside the documented `(n_plans, n_samples)` layout, an empty plan stack and rows that are not permutations are now rejected (PR #911, Issue #910) - Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264) - Fix `ot.bsp.compute_bspot_bijection` (and `random_state` of `ot.solve_sample(..., method="bsp")`) not being reproducible for a given `seed`: its candidate plans came from a `thread_local` `std::mt19937` seeded once per OpenMP thread, so the result depended on the call history and on the number of threads. Every thread started from the same state, so the previous version could underperform with multiple threads. Each plan now draws its own stream, derived from the seed and the plan index (PR #898, Issue #897) - Fix duplicate function and class descriptions in generated module API pages while preserving top-level aliases and example backreferences (Issue #527) diff --git a/ot/bsp/bsp_ot.py b/ot/bsp/bsp_ot.py index ca98c9848..4f689de30 100644 --- a/ot/bsp/bsp_ot.py +++ b/ot/bsp/bsp_ot.py @@ -170,10 +170,8 @@ def merge_bijections(X, Y, perms, p=2): if perms.shape[0] == 0: raise ValueError("perms must contain at least one plan") - # Check the entries before casting to int32: the cast truncates fractional values and - # wraps integers that do not fit, either of which can turn an invalid row into a valid - # one. The merge indexes the point clouds with these values, so they must be a - # permutation of 0..n_samples-1 first. + # Checked before the cast to int32: that cast truncates fractional values and wraps + # integers that do not fit, either of which can turn an invalid row into a permutation. expected = np.broadcast_to(np.arange(X_np.shape[0]), perms.shape) if not np.array_equal(np.sort(perms, axis=1), expected): raise ValueError("each row of perms must be a permutation of 0..n_samples-1") diff --git a/ot/bsp/bsp_wrap.pyx b/ot/bsp/bsp_wrap.pyx index 38c2e6448..23b4830b1 100644 --- a/ot/bsp/bsp_wrap.pyx +++ b/ot/bsp/bsp_wrap.pyx @@ -90,24 +90,6 @@ def merge_bijections_c(np.ndarray[double, ndim=2, mode="c"] X, np.ndarray[double if k == 0: raise ValueError("plans must contain at least one plan") - # C++ indexes the point clouds with these entries, so a stack whose rows are not a - # permutation of 0..n-1 must not reach BijectiveMatching. Counting is linear in the - # size of the stack and does not copy it. - cdef np.ndarray[int, ndim=1, mode="c"] seen = np.zeros(n, dtype=np.int32) - cdef int i - cdef int j - cdef int v - - for i in range(k): - for j in range(n): - seen[j] = 0 - for j in range(n): - v = plans[i, j] - if v < 0 or v >= n or seen[v] != 0: - raise ValueError( - "each row of plans must be a permutation of 0..n_samples-1") - seen[v] = 1 - cdef np.ndarray[int, ndim=1, mode="c"] plan = np.zeros(n, dtype=np.int32) cdef double cost_val diff --git a/test/test_bspot.py b/test/test_bspot.py index 950875154..9c939acd7 100644 --- a/test/test_bspot.py +++ b/test/test_bspot.py @@ -145,23 +145,25 @@ def test_bsp_ot_torch_backend(): def _mean_squared_cost(A, B, perm): - # cost of ``perm`` computed here, independently of what the solver reports + # cost of ``perm`` recomputed here, independently of what the solver reports return float(np.mean(np.sum((A - B[np.asarray(perm)]) ** 2, axis=1))) +def _clouds(n=8, data_seed=0): + rng = np.random.RandomState(data_seed) + return rng.randn(n, 2), rng.randn(n, 2) + + def test_bsp_merge_fewer_plans_than_samples(): - # perms is a stack of (n_plans, n_samples) permutations, so the merge has n_plans of - # them to read. With one plan over 32 samples the old code read one plan per sample - # and walked off the end of the array (issue #910). + # perms stacks (n_plans, n_samples) permutations, so the merge has n_plans of them to + # read. With one plan over 32 samples the old code read one plan per sample and walked + # off the end of the array (issue #910). n = 32 - rng = np.random.RandomState(0) - - A = rng.randn(n, 2) - B = rng.randn(n, 2) + A, B = _clouds(n) perms = np.arange(n, dtype=np.int32).reshape(1, n) with warnings.catch_warnings(): - # merge_bijections warns when its reported cost does not match the permutation + # merge_bijections warns when its reported cost is not the cost of the permutation warnings.simplefilter("error") cost, perm = ot.bsp.merge_bijections(A, B, perms, 2) @@ -172,13 +174,10 @@ def test_bsp_merge_fewer_plans_than_samples(): def test_bsp_merge_uses_every_supplied_plan(): - # The plans past the n_samples-th one must reach the merge. Both clouds are the same, - # so the identity is the only plan with zero cost and the only one that can produce - # it: if it is dropped the merged cost stays on the mediocre plan. + # The plans past the n_samples-th must reach the merge. Both clouds are the same, so + # the identity is the only plan that can give a zero cost. n = 8 - rng = np.random.RandomState(0) - - A = rng.randn(n, 2) + A, _ = _clouds(n) B = A.copy() mediocre = np.roll(np.arange(n), 1).astype(np.int32) @@ -193,46 +192,11 @@ def test_bsp_merge_uses_every_supplied_plan(): np.testing.assert_allclose(cost_with, 0.0, atol=1e-12) -def test_bsp_merge_rejects_plans_of_wrong_length(): - # A stack whose rows are not one entry per point cannot be merged; it must fail here - # instead of being read as a number of plans by C++. - n = 8 - rng = np.random.RandomState(0) - - A = rng.randn(n, 2) - B = rng.randn(n, 2) - - with pytest.raises(ValueError, match="n_samples is the number of points"): - ot.bsp.merge_bijections(A, B, np.tile(np.roll(np.arange(7), 1), (3, 1)), 2) - - with pytest.raises(ValueError, match="2D"): - ot.bsp.merge_bijections(A, B, np.arange(n, dtype=np.int32), 2) - - with pytest.raises(ValueError, match="2D"): - ot.bsp.merge_bijections(A, B, np.zeros((n, n, 1), dtype=np.int32), 2) - - # the same refusal at the native boundary, so no caller can reach the overrun - from ot.bsp import bsp_wrap - - with pytest.raises(ValueError, match="one entry per point"): - bsp_wrap.merge_bijections_c( - np.ascontiguousarray(A, dtype=np.float64), - np.ascontiguousarray(B, dtype=np.float64), - np.ascontiguousarray( - np.tile(np.roll(np.arange(7), 1), (3, 1)), dtype=np.int32 - ), - 2, - ) - - def test_bsp_merge_on_a_standard_stack(): - # The ordinary case must keep its documented behaviour: a bijection no worse than any - # plan it merged, and a cost that is the cost of that bijection. + # The ordinary case keeps its documented behaviour: a bijection no worse than any plan + # it merged, and a cost that is the cost of that bijection. n = 50 - rng = np.random.RandomState(0) - - A = rng.randn(n, 2) - B = rng.randn(n, 2) + A, B = _clouds(n) perms = np.stack([np.roll(np.arange(n), s) for s in (1, 7, 23, 41)]).astype(np.int32) with warnings.catch_warnings(): @@ -246,14 +210,11 @@ def test_bsp_merge_on_a_standard_stack(): assert cost <= _mean_squared_cost(A, B, plan) + 1e-12 -def test_bsp_merge_converts_dtypes_and_refuses_views(): - # np.asarray(perms, dtype=np.int32) already accepted these, and the fix must not change - # that; a non-contiguous view was refused before the fix too, rather than read loosely. +def test_bsp_merge_supports_documented_dtypes_and_refuses_views(): + # int64, float64 and list input were accepted before the fix and must stay so; a + # non-contiguous view was already refused rather than read. n = 8 - rng = np.random.RandomState(0) - - A = rng.randn(n, 2) - B = rng.randn(n, 2) + A, B = _clouds(n) plans = np.tile(np.roll(np.arange(n), 1), (4, 1)) ref_cost, ref_perm = ot.bsp.merge_bijections(A, B, plans.astype(np.int32), 2) @@ -268,100 +229,66 @@ def test_bsp_merge_converts_dtypes_and_refuses_views(): ot.bsp.merge_bijections(A, B, view, 2) -def test_bsp_merge_rejects_empty_plan_stack(): - # An empty stack leaves the merge helper with nothing to read, which crashes the - # interpreter, so it has to be refused before the helper is reached. - n = 8 - rng = np.random.RandomState(0) - - A = rng.randn(n, 2) - B = rng.randn(n, 2) - - with pytest.raises(ValueError, match="at least one plan"): - ot.bsp.merge_bijections(A, B, np.zeros((0, n), dtype=np.int32), 2) - - # the native entry point is reachable on its own, so it must refuse it too +@pytest.mark.parametrize( + "perms, reason", + [ + pytest.param(np.arange(8, dtype=np.int32), "2D", id="one_dimensional"), + pytest.param(np.zeros((8, 8, 1), np.int32), "2D", id="three_dimensional"), + pytest.param( + np.tile(np.roll(np.arange(7), 1), (3, 1)), + "n_samples is the number of points", + id="rows_of_wrong_length", + ), + pytest.param(np.zeros((0, 8), np.int32), "at least one plan", id="empty_stack"), + pytest.param([[0, 0, 2, 3, 4, 5, 6, 7]], "permutation of", id="repeated_entry"), + pytest.param([[-1, 1, 2, 3, 4, 5, 6, 7]], "permutation of", id="negative_index"), + pytest.param([[0, 1, 2, 3, 4, 5, 6, 8]], "permutation of", id="index_past_last_point"), + pytest.param([[0, 1, 2, 3, 4, 5, 6, 6]], "permutation of", id="missing_entry"), + pytest.param([[0, 1, 2, np.nan, 4, 5, 6, 7]], "permutation of", id="nan_entry"), + ], +) +def test_bsp_merge_rejects_invalid_perms(perms, reason): + # A row that is not a permutation of 0..n_samples-1 makes the merge index outside the + # point clouds, and a stack that is not (n_plans, n_samples) cannot be read at all. + A, B = _clouds() + with pytest.raises(ValueError, match=reason): + ot.bsp.merge_bijections(A, B, perms, 2) + + +@pytest.mark.parametrize( + "perms", + [ + pytest.param(np.array([0, 1, 2, 3.5, 4, 5, 6, 7]), id="truncates_to_valid"), + pytest.param( + np.array([0, 1, 2, 2**32 + 3, 4, 5, 6, 7], dtype=np.int64), + id="wraps_to_valid", + ), + ], +) +def test_bsp_merge_checks_the_entries_before_casting_to_int32(perms): + # Reading these rows as int32 yields a valid permutation -- the cast truncates 3.5 and + # keeps the low bits of 2**32 + 3 -- so only a check made before the cast refuses them. + A, B = _clouds() + np.testing.assert_array_equal(perms.astype(np.int32), np.arange(8)) + with pytest.raises(ValueError, match="permutation of"): + ot.bsp.merge_bijections(A, B, perms.reshape(1, 8), 2) + + +def test_bsp_merge_c_refuses_a_stack_it_cannot_bound(): + # merge_bijections_c is reachable on its own, so it keeps the row-length and the + # empty-stack guards that bound the buffer it hands to C++. from ot.bsp import bsp_wrap - with pytest.raises(ValueError, match="at least one plan"): - bsp_wrap.merge_bijections_c( - np.ascontiguousarray(A), - np.ascontiguousarray(B), - np.zeros((0, n), dtype=np.int32), - 2, - ) - - -def test_bsp_merge_rejects_rows_that_are_not_permutations(): - # A row describes a bijection only if it is a permutation of 0..n-1. Anything else - # makes the merge index outside the point clouds, so it is refused rather than read. - n = 8 - rng = np.random.RandomState(0) - - A = rng.randn(n, 2) - B = rng.randn(n, 2) - identity = np.arange(n, dtype=np.int32) - - invalid = [ - np.array([0, 0, 2, 3, 4, 5, 6, 7], dtype=np.int32), # repeated entry - np.array([-1, 1, 2, 3, 4, 5, 6, 7], dtype=np.int32), # negative index - np.array([0, 1, 2, 3, 4, 5, 6, 8], dtype=np.int32), # past the last point - np.array([0, 1, 2, 3, 4, 5, 6, 6], dtype=np.int32), # missing entry - ] - - from ot.bsp import bsp_wrap - - for row in invalid: - stack = row.reshape(1, n) - with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): - ot.bsp.merge_bijections(A, B, stack, 2) - with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): + A, B = _clouds() + for perms, reason in [ + (np.tile(np.roll(np.arange(7), 1), (3, 1)), "one entry per point"), + (np.zeros((0, 8), np.int32), "at least one plan"), + ]: + with pytest.raises(ValueError, match=reason): bsp_wrap.merge_bijections_c( - np.ascontiguousarray(A), - np.ascontiguousarray(B), - np.ascontiguousarray(stack, dtype=np.int32), - 2, + A, B, np.ascontiguousarray(perms, dtype=np.int32), 2 ) - # one bad row is enough to refuse the whole stack, wherever it sits in it - mixed = np.stack([identity, row, identity]) - with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): - ot.bsp.merge_bijections(A, B, mixed, 2) - with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): - bsp_wrap.merge_bijections_c( - np.ascontiguousarray(A), - np.ascontiguousarray(B), - np.ascontiguousarray(mixed, dtype=np.int32), - 2, - ) - - -def test_bsp_merge_rejects_plans_that_only_look_valid_after_casting(): - # The entries are checked before the conversion to int32. That cast truncates - # fractional values and wraps integers that do not fit, and either one would silently - # turn an invalid row into a permutation if the check ran after it. - n = 8 - rng = np.random.RandomState(0) - - A = rng.randn(n, 2) - B = rng.randn(n, 2) - - fractional = np.arange(n, dtype=np.float64).reshape(1, n) - fractional[0, 3] = 3.5 - with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): - ot.bsp.merge_bijections(A, B, fractional, 2) - - wrapped = np.arange(n, dtype=np.int64).reshape(1, n) - wrapped[0, 3] = 2**32 + 3 - assert wrapped.astype(np.int32)[0, 3] == 3, "the cast keeps the low bits only" - with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): - ot.bsp.merge_bijections(A, B, wrapped, 2) - - nan = np.arange(n, dtype=np.float64).reshape(1, n) - nan[0, 5] = np.nan - with pytest.raises(ValueError, match="permutation of 0..n_samples-1"): - ot.bsp.merge_bijections(A, B, nan, 2) - def _two_clouds(n=100, data_seed=0): rng = np.random.RandomState(data_seed)