Merge pull request #16551 from numpy/fix/seed-sequence-zeros

BUG: Ensure SeedSequence 0-padding does not collide with spawn keys

2 files changed, 37 insertions, 2 deletions

diff --git a/numpy/random/bit_generator.pyx b/numpy/random/bit_generator.pyx
index f145ec13d..3c52a9933 100644
--- a/numpy/random/bit_generator.pyx
+++ b/numpy/random/bit_generator.pyx
@@ -382,13 +382,22 @@ cdef class SeedSequence():
         -------
         entropy_array : 1D uint32 array
         """
-        # Convert run-entropy, program-entropy, and the spawn key into uint32
+        # Convert run-entropy and the spawn key into uint32
         # arrays and concatenate them.
 
         # We MUST have at least some run-entropy. The others are optional.
         assert self.entropy is not None
         run_entropy = _coerce_to_uint32_array(self.entropy)
         spawn_entropy = _coerce_to_uint32_array(self.spawn_key)
+        if len(spawn_entropy) > 0 and len(run_entropy) < self.pool_size:
+            # Explicitly fill out the entropy with 0s to the pool size to avoid
+            # conflict with spawn keys. We changed this in 1.19.0 to fix
+            # gh-16539. In order to preserve stream-compatibility with
+            # unspawned SeedSequences with small entropy inputs, we only do
+            # this when a spawn_key is specified.
+            diff = self.pool_size - len(run_entropy)
+            run_entropy = np.concatenate(
+                [run_entropy, np.zeros(diff, dtype=np.uint32)])
         entropy_array = np.concatenate([run_entropy, spawn_entropy])
         return entropy_array
 
diff --git a/numpy/random/tests/test_seed_sequence.py b/numpy/random/tests/test_seed_sequence.py
index fe23680ed..f08cf80fa 100644
--- a/numpy/random/tests/test_seed_sequence.py
+++ b/numpy/random/tests/test_seed_sequence.py
@@ -1,5 +1,5 @@
 import numpy as np
-from numpy.testing import assert_array_equal
+from numpy.testing import assert_array_equal, assert_array_compare
 
 from numpy.random import SeedSequence
 
@@ -52,3 +52,29 @@ def test_reference_data():
         assert_array_equal(state, expected)
         state64 = ss.generate_state(len(expected64), dtype=np.uint64)
         assert_array_equal(state64, expected64)
+
+
+def test_zero_padding():
+    """ Ensure that the implicit zero-padding does not cause problems.
+    """
+    # Ensure that large integers are inserted in little-endian fashion to avoid
+    # trailing 0s.
+    ss0 = SeedSequence(42)
+    ss1 = SeedSequence(42 << 32)
+    assert_array_compare(
+        np.not_equal,
+        ss0.generate_state(4),
+        ss1.generate_state(4))
+
+    # Ensure backwards compatibility with the original 0.17 release for small
+    # integers and no spawn key.
+    expected42 = np.array([3444837047, 2669555309, 2046530742, 3581440988],
+                          dtype=np.uint32)
+    assert_array_equal(SeedSequence(42).generate_state(4), expected42)
+
+    # Regression test for gh-16539 to ensure that the implicit 0s don't
+    # conflict with spawn keys.
+    assert_array_compare(
+        np.not_equal,
+        SeedSequence(42, spawn_key=(0,)).generate_state(4),
+        expected42)