state.py

1import builtins2import numpy as np3from typing import Any, Optional45class State:67    def __init__(self):8        self._rng = np.random.RandomState()910    @property11    def rng(self):12        return self._rng1314    def seed(self, seed: Optional[int] = None):15        self._rng = np.random.RandomState(seed)1617state = State()1819# API2021def seed(s: Optional[int] = None):22    state.seed(s)2324def bool() -> builtins.bool:25    return builtins.bool(state.rng.choice([True, False]))2627def choice(seq: list) -> Any:28    idx = state.rng.randint(0, len(seq))29    return seq[idx]3031def sample(seq: list, k: int) -> list:32    indices = state.rng.choice(len(seq), size=k, replace=False)33    return [seq[i] for i in indices]3435def shuffle(seq: list):36    state.rng.shuffle(seq)3738def randint(a: int, b: int) -> int:39    return int(state.rng.randint(a, b + 1))