add step-setup

pymc/sampling/mcmc.py

@@ -1527,7 +1527,7 @@ def _iter_sample(
     if draws < 1:
         raise ValueError("Argument `draws` must be greater than 0.")
 
-    step.set_rng(rng)
+    step.setup_chain(rng, tune, draws - tune)
 
     point = start
     if isinstance(trace, _ZarrChainBase):

pymc/sampling/parallel.py

@@ -248,7 +248,7 @@ def _recv_msg(self):
 
     def _start_loop(self):
         zarr_recording = self._zarr_recording
-        self._step_method.set_rng(self._rng)
+        self._step_method.setup_chain(self._rng, self._tune, self._draws - self._tune)
 
         draw = 0
         tuning = True

pymc/sampling/population.py

@@ -396,7 +396,7 @@ def _prepare_iter_population(
             chainstep = CompoundStep([copy(m) for m in step.methods])
         else:
             chainstep = copy(step)
-        chainstep.set_rng(rng)
+        chainstep.setup_chain(rng, tune, draws - tune)
         # link population samplers to the shared population state
         for sm in chainstep.methods if isinstance(step, CompoundStep) else [chainstep]:
             if isinstance(sm, PopulationArrayStepShared):

pymc/step_methods/compound.py

@@ -230,7 +230,13 @@ def stop_tuning(self):
         if hasattr(self, "tune"):
             self.tune = False
 
-    def set_rng(self, rng: RandomGenerator):
+    def setup_chain(self, rng: RandomGenerator, tune: int, draws: int) -> None:
+        """Prepare the step method for sampling one chain.
+
+        Called once per chain, right before that chain starts sampling.
+        Can be called multiple times, meaning downstream method
+        should handle multiple, possibly repeated calls.
+        """
         self.rng = get_random_generator(rng, copy=False)
 
 
@@ -297,6 +303,11 @@ def reset_tuning(self):
             if hasattr(method, "reset_tuning"):
                 method.reset_tuning()
 
+    def setup_chain(self, rng: RandomGenerator, tune: int, draws: int) -> None:
+        _rngs = get_random_generator(rng, copy=False).spawn(len(self.methods))
+        for method, _rng in zip(self.methods, _rngs):
+            method.setup_chain(_rng, tune, draws)
+
     @property
     def sampling_state(self) -> DataClassState:
         return CompoundStepState(methods=[method.sampling_state for method in self.methods])
@@ -313,11 +324,6 @@ def sampling_state(self, state: DataClassState):
     def vars(self) -> list[Variable]:
         return [var for method in self.methods for var in method.vars]
 
-    def set_rng(self, rng: RandomGenerator):
-        _rngs = get_random_generator(rng, copy=False).spawn(len(self.methods))
-        for method, _rng in zip(self.methods, _rngs):
-            method.set_rng(_rng)
-
     def _progressbar_config(self, n_chains=1):
         from functools import reduce
 

pymc/step_methods/hmc/base_hmc.py

@@ -39,7 +39,7 @@
 from pymc.step_methods.state import dataclass_state
 from pymc.step_methods.step_sizes import DualAverageAdaptation, StepSizeState
 from pymc.tuning import guess_scaling
-from pymc.util import RandomGenerator, get_random_generator, get_value_vars_from_user_vars
+from pymc.util import RandomGenerator, get_value_vars_from_user_vars
 
 logger = logging.getLogger(__name__)
 
@@ -297,6 +297,6 @@ def reset(self, start=None):
         self.tune = True
         self.potential.reset()
 
-    def set_rng(self, rng: RandomGenerator):
-        self.rng = get_random_generator(rng, copy=False)
+    def setup_chain(self, rng: RandomGenerator, tune: int, draws: int) -> None:
+        super().setup_chain(rng, tune, draws)
         self.potential.set_rng(self.rng.spawn(1)[0])