jaxvacua.util.random_integer

Contents

jaxvacua.util.random_integer#

random_integer(lower_bound, upper_bound, rns_key=None, seed=42, shape=(1,))#

Sample uniformly distributed integers on [lower_bound, upper_bound] (inclusive on both ends, matching the convention of jax.random.randint(maxval=upper_bound + 1)).

Parameters:
  • lower_bound (int) – Lower edge of the sampling interval.

  • upper_bound (int) – Upper edge of the sampling interval (inclusive).

  • rns_key (Union[PRNGSequence, Array, None]) – Key source — see random_uniform() for details.

  • seed (int) – Seed used when rns_key is None. Default 42.

  • shape (Tuple[int, ...]) – Output shape. Default (1,).

Returns:

Array – Integer samples of shape shape.

Return type:

Array