Module random 

Random numbers from explicit keys or a per-thread key sequence, bit for bit as in MLX.

A key is a uint32 array of shape [2]. bits hashes counters with threefry2x32 using MLX’s layout, and uniform and normal are composed from operations exactly as in MLX’s random.cpp. For the same seed, bits and uniform samples equal MLX’s. Normal samples agree within a few ulps: erfinv needs a logarithm, whose last bits differ between implementations.

Without a key, functions draw the next key from this thread’s sequence. Each thread’s sequence starts from a key derived from the process seed (a time seed until seed is called) and the thread’s ordinal, the order in which threads first draw:

• ordinal 0 starts from key(process_seed), so a single-threaded program seeded with seed(s) draws exactly MLX’s values for s;
• ordinal n > 0 starts from key(splitmix64(process_seed ^ n)).

So threads that never called seed draw distinct streams. seed sets the process seed, restarts the calling thread’s sequence as ordinal 0 and makes every other thread derive its sequence again (with the next ordinals, in the order they draw next); a thread that called seed itself keeps its own sequence until it calls seed again. Draws are reproducible across runs when threads start drawing in a fixed order; for exact reproducibility of worker threads whatever their order, pass explicit keys (key, split, split_n).

Like every operation, these return arrays: invalid arguments give a failed array that carries the error and the caller’s location (see crate::Array::check).

Functions

bernoulli

Booleans of shape that are true with probability p, as MLX: uniform < p.

bits

Random bits of shape, width bytes per element (1, 2 or 4 → uint8, uint16, uint32).

categorical

Indices (uint32) drawn from the categorical distributions whose unnormalised log probabilities run along one axis of logits (..: the flattened array), as MLX: the argmax of logits plus Gumbel noise. The result has logits’ shape without axis.

gumbel

Gumbel samples of shape and floating-point dtype, as MLX: −log(−log(u)) for u uniform in [0, 1).

key

The body of crate::random::key.

next_key

The next key of this thread’s sequence (split the current key; keep the first half), the key that functions called without one use. Useful to pass the sequence’s state into a compiled function, where keys drawn inside would be frozen at the first trace.

normal

Normal samples (mean loc, standard deviation scale; MLX’s defaults are 0 and 1) of shape and floating-point dtype, as MLX: √2 · erfinv(u) for u uniform in (-1, 1).

randint

Integers in [low, high) of shape and integer dtype, as MLX: uniform floats in [low, high) rounded down.

seed

The body of crate::random::seed.

split

num new keys from key, as an array of shape [num, 2].

split_n

N new keys from key, for destructuring: let [k1, k2] = random::split_n::<2>(&key);.

uniform

Uniform samples in [low, high) of shape and floating-point dtype (float16, bfloat16 or float32), as MLX: 32 random bits divided by u32::MAX, clamped below 1.