add rng in NelderMeadSampler argment.

aistairc · Mar 26, 2024 · 32bf9d3 · 32bf9d3
1 parent cc82cd1
commit 32bf9d3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/aiaccel/hpo/samplers/nelder_mead_sampler.py b/aiaccel/hpo/samplers/nelder_mead_sampler.py
@@ -241,15 +241,17 @@ def __init__(
         self,
         search_space: dict[str, tuple[float, float]],
         seed: int | None = None,
+        rng: np.random.RandomState | None = None,
         coeff: NelderMeadCoefficient | None = None,
         parallel_enabled: bool = False,
     ) -> None:
         self._search_space = search_space
+        _rng = rng if rng is not None else np.random.RandomState(seed) if seed is not None else np.random.RandomState()
 
         self.nm = NelderMeadAlgorism(
             search_space=self._search_space,
             coeff=coeff,
-            rng=np.random.RandomState(seed),
+            rng=_rng,
             block=parallel_enabled,
             timeout=None
             )