Random Number Generation

rng ¶

A sub-package to handle random number generators.

RNGManager(seed: int | None = None)

Manager for random number generation with state persistence.

Initialize the RNG manager.

Parameters:

Name	Type	Description	Default
`seed`	`int \| None`	Random seed. If None, draw a random integer between [0, 2**63).	`None`

key: Array

Get the random number generator key.

Returns:

Type	Description
`Array`	JAX random key.

key_data: Array

Get the key data.

The function can be used to retrieve the key data for storing in disk.

Returns:

Type	Description
`Array`	Key data.

new_key: Array

Get a new key.

Returns:

Type	Description
`Array`	A new key.

save_key(path: str | Path) -> None

Save RNG key to file.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Path to save the key file.	required

load_key(path: str | Path) -> None

Load RNG key from file.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Path to load the key file.	required