Module `Rune.Rng`

Explicit splittable PRNG keys.

Randomness is a pure function of a key: a sampler applied to the same key, dtype, and shape returns the same values — eagerly, under jit, and on every device. Rng.uniform, Rng.randint and Rng.bernoulli are bit-identical everywhere; Rng.normal agrees up to the device's transcendental rounding. A key is an ordinary [|2|] int32 tensor, so it goes wherever tensors go: a leaf of a parameter structure, an input of a jitted function, a mapped axis of vmap.

Fresh randomness comes from fresh keys, never from calling a sampler twice: Rng.split a key into independent subkeys and hand one to each consumer, or Rng.fold_in a step counter to derive per-step keys from a root.

let root = Rune.Rng.key 42 in
let step = Rune.jit (module P) (fun p -> ... Rune.Rng.uniform p.key ...) in
for i = 0 to steps - 1 do
  ignore (step { params with key = Rune.Rng.fold_in root i })
done

type key = (int32, Nx.int32_elt) Nx.t

The type for PRNG keys: an [|2|] int32 tensor holding the generator's state. Keys are values — thread them explicitly, and never reuse one for two draws that must be independent.

val key : int -> key

key seed is the key for seed. Equal seeds give equal keys.

val split : ?n:int -> key -> key array

split ?n k is n independent subkeys derived from k (default 2). Deterministic: splitting the same key yields the same subkeys. The subkeys are statistically independent of each other; derive one subkey per consumer instead of reusing k.

val fold_in : key -> int -> key

fold_in k data is the subkey of k indexed by data: distinct data values give independent keys. Use it to derive per-step keys from a root key and a loop counter.

val uniform : key -> (float, 'b) Nx.dtype -> int array -> (float, 'b) Nx.t

uniform k dtype shape samples uniformly from [0, 1). Pure: the same key, dtype and shape always produce the same values.

val normal : key -> (float, 'b) Nx.dtype -> int array -> (float, 'b) Nx.t

normal k dtype shape samples the standard normal distribution (Box-Muller over two uniform draws).

val randint : 
  key ->
  ?high:int ->
  int array ->
  int ->
  (int32, Nx.int32_elt) Nx.t

randint k ?high shape low samples integers uniformly from [low, high). high defaults to 10.

Raises Invalid_argument if low >= high.

val bernoulli : key -> p:float -> int array -> (bool, Nx.bool_elt) Nx.t

bernoulli k ~p shape samples booleans that are true with probability p.

Raises Invalid_argument if p is outside [0, 1].