PrincetonUniversity / PsyNeuLink / 25535066478

84%

master: 85%

Build Type

push

github

Committed by web-flow

Commit Message

OneHot: fix PROB sampling at the cumsum-rounding boundary (#3535)

## Summary

`OneHot` `PROB`/`PROB_INDICATOR` sampling could fail when
`np.cumsum(prob_dist)[-1]` drifted just below 1.0 due to float rounding
(e.g. `np.cumsum([0.7, 0.2, 0.1])[-1] == 0.9999999999999999`) and
`random_state.uniform()` returned a value in the resulting
`(cum_sum[-1], 1.0)` gap.

* **Python path** (`selectionfunctions.py:733`) raised `StopIteration`
from the unguarded generator
  ```python
chosen_item = next(element for element in cum_sum if element >
random_value)
  ```
* **LLVM path** (`_gen_llvm_function_body` for `PROB`/`PROB_INDICATOR`)
silently produced an all-zero output, because no iteration's `sum_old <=
random_draw < sum_new` condition fired.

The Python path now uses `np.searchsorted(..., side='right')` clamped to
the last index, and indexes the chosen position directly. This also
fixes a second-order bug where `np.where(cum_sum == chosen_item, 1, 0)`
selected multiple indices when `prob_dist` contained zeros (which
produced duplicate `cum_sum` entries) — see the corrected `SOFT_MAX
MASK_THRESHOLD PROB` expectations in `test_transfer.py`.

The LLVM path forces the upper-bound check to fire on the last
iteration, ensuring the last index always wins when the draw exceeds the
final prefix sum.

## Tests

* `test_one_hot_prob_cumsum_below_one_due_to_float_rounding` —
deterministic Python regression test that stubs `random_state.uniform()`
to return a value in the gap. Reproduces the `StopIteration` against the
unfixed code.
* `test_one_hot_prob_runtime_cumsum_under_one_never_all_zero` — covers
both Python and LLVM paths by constructing `OneHot` with a valid default
`prob_dist` (so validation passes) and then calling it with a runtime
`prob_dist` whose `cum_sum[-1] < 1.0`. About half of all uniform draws
then land in the gap, deterministically triggering both bugs against the
unfixed code (Python: `StopIteration`; LLVM: all-zero output).

## Test plan

- [x] `pytest ... (continued)

Coverage Stats

10519 of 13698 branches covered (76.79%)

Branch coverage included in aggregate %.

6 of 6 new or added lines in 1 file covered. (100.0%)

3 existing lines in 2 files now uncovered.

36185 of 41762 relevant lines covered (86.65%)

0.87 hits per line