DataClassState uses RandomGeneratorState instead of Generator objects

Author: lucianopaz

pymc/step_methods/compound.py

@@ -31,7 +31,12 @@
 
 from pymc.blocking import PointType, StatDtype, StatsDict, StatShape, StatsType
 from pymc.model import modelcontext
-from pymc.step_methods.state import DataClassState, WithSamplingState, dataclass_state
+from pymc.step_methods.state import (
+    DataClassState,
+    RandomGeneratorState,
+    WithSamplingState,
+    dataclass_state,
+)
 from pymc.util import RandomGenerator, get_random_generator
 
 __all__ = ("Competence", "CompoundStep")
@@ -91,7 +96,7 @@ def infer_warn_stats_info(
 
 @dataclass_state
 class StepMethodState(DataClassState):
-    rng: np.random.Generator
+    rng: RandomGeneratorState
 
 
 class BlockedStep(ABC, WithSamplingState):

pymc/step_methods/hmc/quadpotential.py

@@ -26,7 +26,12 @@
 from scipy.sparse import issparse
 
 from pymc.pytensorf import floatX
-from pymc.step_methods.state import DataClassState, WithSamplingState, dataclass_state
+from pymc.step_methods.state import (
+    DataClassState,
+    RandomGeneratorState,
+    WithSamplingState,
+    dataclass_state,
+)
 from pymc.util import RandomGenerator, get_random_generator
 
 __all__ = [
@@ -105,7 +110,7 @@ def __str__(self):
 
 @dataclass_state
 class PotentialState(DataClassState):
-    rng: np.random.Generator
+    rng: RandomGeneratorState
 
 
 class QuadPotential(WithSamplingState):

pymc/step_methods/state.py

@@ -17,6 +17,8 @@
 
 import numpy as np
 
+from pymc.util import RandomGeneratorState, get_state_from_generator, random_generator_from_state
+
 dataclass_state = dataclass(kw_only=True)
 
 
@@ -66,8 +68,11 @@ def sampling_state(self) -> DataClassState:
         kwargs = {}
         for field in fields(state_class):
             val = getattr(self, field.name)
+            _val: Any
             if isinstance(val, WithSamplingState):
                 _val = val.sampling_state
+            elif isinstance(val, np.random.Generator):
+                _val = get_state_from_generator(val)
             else:
                 _val = val
             kwargs[field.name] = deepcopy(_val)
@@ -81,6 +86,8 @@ def sampling_state(self, state: DataClassState):
         ), f"Encountered invalid state class '{state.__class__}'. State must be '{state_class}'"
         for field in fields(state_class):
             state_val = deepcopy(getattr(state, field.name))
+            if isinstance(state_val, RandomGeneratorState):
+                state_val = random_generator_from_state(state_val)
             self_val = getattr(self, field.name)
             is_frozen = field.metadata.get("frozen", False)
             if is_frozen:

pymc/util.py

@@ -15,6 +15,7 @@
 import functools
 import warnings
 
+from collections import namedtuple
 from collections.abc import Sequence
 from copy import deepcopy
 from typing import NewType, cast
@@ -601,6 +602,31 @@ def update(
         return None
 
 
+RandomGeneratorState = namedtuple("RandomGeneratorState", ["bit_generator_state", "seed_seq_state"])
+
+
+def get_state_from_generator(
+    rng: np.random.Generator | np.random.BitGenerator,
+) -> RandomGeneratorState:
+    assert isinstance(rng, (np.random.Generator | np.random.BitGenerator))
+    bit_gen: np.random.BitGenerator = (
+        rng.bit_generator if isinstance(rng, np.random.Generator) else rng
+    )
+
+    return RandomGeneratorState(
+        bit_generator_state=bit_gen.state,
+        seed_seq_state=bit_gen.seed_seq.state,  # type: ignore[attr-defined]
+    )
+
+
+def random_generator_from_state(state: RandomGeneratorState) -> np.random.Generator:
+    seed_seq = np.random.SeedSequence(**state.seed_seq_state)
+    bit_generator_class = getattr(np.random, state.bit_generator_state["bit_generator"])
+    bit_generator = bit_generator_class(seed_seq)
+    bit_generator.state = state.bit_generator_state
+    return np.random.Generator(bit_generator)
+
+
 def get_random_generator(
     seed: RandomGenerator | np.random.RandomState = None, copy: bool = True
 ) -> np.random.Generator:
@@ -645,6 +671,10 @@ def get_random_generator(
         # In the former case, it will return seed, in the latter it will return
         # a new Generator object that has the same BitGenerator. This would potentially
         # make the new generator be shared across many users. To avoid this, we
-        # deepcopy by default.
+        # copy by default.
+        # Also, because of https://github.com/numpy/numpy/issues/27727, we can't use
+        # deepcopy. We must rebuild a Generator without losing the SeedSequence information
+        if isinstance(seed, np.random.Generator | np.random.BitGenerator):
+            return random_generator_from_state(get_state_from_generator(seed))
         seed = deepcopy(seed)
     return np.random.default_rng(seed)

tests/step_methods/test_metropolis.py

@@ -14,8 +14,6 @@
 
 import warnings
 
-from copy import deepcopy
-
 import arviz as az
 import numpy as np
 import numpy.testing as npt
@@ -406,8 +404,7 @@ def test_sampling_state(step_method, model_fn):
             sampler = step_method(model.value_vars)
             if hasattr(sampler, "link_population"):
                 sampler.link_population([initial_point] * 100, 0)
-            sampler_orig = deepcopy(sampler)
-            state_orig = sampler_orig.sampling_state
+            state_orig = sampler.sampling_state
 
             sample1, stat1 = sampler.step(initial_point)
             sampler.tune = False