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, which starts from a per-process time seed until seed is called (as in MLX).

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 axis of logits, 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

A key for seed: [seed >> 32, seed & 0xffffffff].

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) 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

Restarts this thread’s key sequence from seed.

split

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

split_pair

Two new keys from 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.