bilby.core.utils.random.random_array_module

bilby.core.utils.random.random_array_module(random_state)[source]

Return the array module corresponding to the provided random state. The the random state is a JAX random key, this will return jax.numpy. Otherwise, it will return numpy.

Parameters:
random_state: None, int, np.random.Generator, or jax.random.KeyArray

The random state to resolve.

Returns:
array module

The array module corresponding to the provided random state.