opentau.utils.random_utils

Utilities for managing random number generator states.

This module provides functions for serializing, deserializing, saving, and loading random number generator states for Python’s random module, NumPy, and PyTorch. This is essential for reproducibility in training and evaluation.

Functions

`deserialize_numpy_rng_state`(rng_state_dict)	Restore NumPy's RNG state from a dictionary.
`deserialize_python_rng_state`(rng_state_dict)	Restore Python's random module RNG state from a dictionary.
`deserialize_rng_state`(rng_state_dict)	Restore RNG states for random, numpy, and torch from a dictionary.
`deserialize_torch_rng_state`(rng_state_dict)	Restore PyTorch's RNG state from a dictionary.
`get_rng_state`()	Get the current random state for random, numpy, and torch.
`load_rng_state`(save_dir)	Load RNG state from a file in the specified directory.
`save_rng_state`(save_dir)	Save RNG state to a file in the specified directory.
`seeded_context`(seed)	Set the seed when entering a context, and restore the prior random state at exit.
`serialize_numpy_rng_state`()	Serialize NumPy's RNG state to a dictionary.
`serialize_python_rng_state`()	Serialize Python's random module RNG state to a dictionary.
`serialize_rng_state`()	Serialize RNG states for random, numpy, and torch.
`serialize_torch_rng_state`()	Serialize PyTorch's RNG state to a dictionary.
`set_rng_state`(random_state_dict)	Set the random state for random, numpy, and torch.
`set_seed`(seed[, accelerator])	Set seed for reproducibility across random, numpy, and torch.

opentau.utils.random_utils.deserialize_numpy_rng_state(rng_state_dict: dict[str, Tensor]) → None[source]

Restore NumPy’s RNG state from a dictionary.

Parameters:: rng_state_dict – Dictionary produced by serialize_numpy_rng_state().

opentau.utils.random_utils.deserialize_python_rng_state(rng_state_dict: dict[str, Tensor]) → None[source]

Restore Python’s random module RNG state from a dictionary.

Parameters:: rng_state_dict – Dictionary produced by serialize_python_rng_state().

opentau.utils.random_utils.deserialize_rng_state(rng_state_dict: dict[str, Tensor]) → None[source]

Restore RNG states for random, numpy, and torch from a dictionary.

Parameters:: rng_state_dict – Dictionary produced by serialize_rng_state().

opentau.utils.random_utils.deserialize_torch_rng_state(rng_state_dict: dict[str, Tensor]) → None[source]

Restore PyTorch’s RNG state from a dictionary.

Parameters:: rng_state_dict – Dictionary produced by serialize_torch_rng_state().

opentau.utils.random_utils.get_rng_state() → dict[str, Any][source]

Get the current random state for random, numpy, and torch.

Returns:: Dictionary containing the current RNG states for all three generators.

opentau.utils.random_utils.load_rng_state(save_dir: Path) → None[source]

Load RNG state from a file in the specified directory.

Parameters:: save_dir – Directory path containing the RNG state file.

opentau.utils.random_utils.save_rng_state(save_dir: Path) → None[source]

Save RNG state to a file in the specified directory.

Parameters:: save_dir – Directory path where the RNG state file will be saved.

opentau.utils.random_utils.seeded_context(seed: int) → Generator[None, None, None][source]

Set the seed when entering a context, and restore the prior random state at exit.

Example usage:

a = random.random()  # produces some random number
with seeded_context(1337):
    b = random.random()  # produces some other random number
c = random.random()  # produces yet another random number, but the same it would have if we never made `b`

opentau.utils.random_utils.serialize_numpy_rng_state() → dict[str, Tensor][source]

Serialize NumPy’s RNG state to a dictionary.

Returns:: Dictionary containing the RNG state as torch.Tensor values, suitable for saving with safetensors.save_file() or torch.save().

opentau.utils.random_utils.serialize_python_rng_state() → dict[str, Tensor][source]

Serialize Python’s random module RNG state to a dictionary.

Returns:: Dictionary containing the RNG state as torch.Tensor values, suitable for saving with safetensors.save_file() or torch.save().

opentau.utils.random_utils.serialize_rng_state() → dict[str, Tensor][source]

Serialize RNG states for random, numpy, and torch.

Returns:: Dictionary containing all RNG states as torch.Tensor values, suitable for saving with safetensors.save_file() or torch.save().

opentau.utils.random_utils.serialize_torch_rng_state() → dict[str, Tensor][source]

Serialize PyTorch’s RNG state to a dictionary.

Returns:: Dictionary containing the RNG state as torch.Tensor values, including CUDA RNG state if available. Suitable for saving with safetensors.save_file() or torch.save().

opentau.utils.random_utils.set_rng_state(random_state_dict: dict[str, Any]) → None[source]

Set the random state for random, numpy, and torch.

Parameters:: random_state_dict – Dictionary of the form returned by get_rng_state().

opentau.utils.random_utils.set_seed(seed, accelerator: Accelerator | None = None) → None[source]

Set seed for reproducibility across random, numpy, and torch.

Parameters:

seed – Seed value to use. If None, no seeding is performed.
accelerator – Optional Accelerator instance. If provided, each process gets a different seed offset to ensure reproducibility in distributed settings.