|
| 1 | +//! Random number generation. |
| 2 | +//! |
| 3 | +//! MLX's RNG is *functional*: a draw is a pure function of an explicit [`Key`], |
| 4 | +//! so the same key always produces the same values. Passing `None` instead |
| 5 | +//! draws from MLX's implicit global state, which [`seed`] initializes. |
| 6 | +//! |
| 7 | +//! Prefer an explicit key when reproducibility matters. The global state is |
| 8 | +//! process-wide, so two threads (or two `#[test]`s running in parallel) that |
| 9 | +//! draw without a key will interleave and neither is reproducible on its own. |
| 10 | +
|
| 11 | +use mlxcore_sys as sys; |
| 12 | + |
| 13 | +use crate::array::Array; |
| 14 | +use crate::dtype::ArrayElement; |
| 15 | +use crate::error::{self, Result}; |
| 16 | +use crate::ffi::{absent_array, as_ffi_ptr}; |
| 17 | +use crate::stream::Stream; |
| 18 | + |
| 19 | +/// A key identifying a position in MLX's random stream. |
| 20 | +/// |
| 21 | +/// Draws made with the same key are identical. To get independent draws, either |
| 22 | +/// build keys from different seeds or [`split`](Self::split) an existing key. |
| 23 | +#[derive(Debug)] |
| 24 | +pub struct Key(Array); |
| 25 | + |
| 26 | +impl Key { |
| 27 | + /// Builds a key from a seed. |
| 28 | + /// |
| 29 | + /// Deterministic: the same seed always yields the same key, and therefore |
| 30 | + /// the same draws. |
| 31 | + pub fn new(seed: u64) -> Result<Self> { |
| 32 | + error::install(); |
| 33 | + let mut out = unsafe { sys::mlx_array_new() }; |
| 34 | + // SAFETY: the result is written into `out`. This takes no stream. |
| 35 | + let status = unsafe { sys::mlx_random_key(&mut out, seed) }; |
| 36 | + Array::from_op(out, status).map(Key) |
| 37 | + } |
| 38 | + |
| 39 | + /// Splits this key into two independent keys. |
| 40 | + /// |
| 41 | + /// Use this to draw more than once from a single seed without reusing a key |
| 42 | + /// (which would give identical values). Splitting is itself deterministic: |
| 43 | + /// the same key always splits into the same pair. |
| 44 | + pub fn split(&self, stream: &Stream) -> Result<(Key, Key)> { |
| 45 | + error::install(); |
| 46 | + let mut first = unsafe { sys::mlx_array_new() }; |
| 47 | + let mut second = unsafe { sys::mlx_array_new() }; |
| 48 | + // SAFETY: both out-params are fresh handles; the key and stream are |
| 49 | + // valid; mlx writes one key into each. |
| 50 | + let status = unsafe { |
| 51 | + sys::mlx_random_split(&mut first, &mut second, self.0.as_raw(), stream.as_raw()) |
| 52 | + }; |
| 53 | + match error::check(status) { |
| 54 | + // SAFETY: on success both handles are owned by us. |
| 55 | + Ok(()) => Ok(unsafe { (Key(Array::from_raw(first)), Key(Array::from_raw(second))) }), |
| 56 | + Err(e) => { |
| 57 | + // SAFETY: both handles were created by mlx and are owned here; |
| 58 | + // free them so the failed call doesn't leak. |
| 59 | + unsafe { |
| 60 | + sys::mlx_array_free(first); |
| 61 | + sys::mlx_array_free(second); |
| 62 | + } |
| 63 | + Err(e) |
| 64 | + } |
| 65 | + } |
| 66 | + } |
| 67 | + |
| 68 | + /// The raw handle for an optional key argument, or an absent handle. |
| 69 | + fn handle(key: Option<&Key>) -> sys::mlx_array { |
| 70 | + match key { |
| 71 | + Some(k) => k.0.as_raw(), |
| 72 | + None => absent_array(), |
| 73 | + } |
| 74 | + } |
| 75 | +} |
| 76 | + |
| 77 | +/// Seeds MLX's implicit global random state. |
| 78 | +/// |
| 79 | +/// Only affects draws made with `key: None`. Because the state is process-wide |
| 80 | +/// and mutated by every unkeyed draw, this does not make a draw reproducible in |
| 81 | +/// the presence of other threads — use a [`Key`] for that. |
| 82 | +pub fn seed(seed: u64) -> Result<()> { |
| 83 | + error::install(); |
| 84 | + // SAFETY: takes a plain integer and mutates global mlx state. |
| 85 | + let status = unsafe { sys::mlx_random_seed(seed) }; |
| 86 | + error::check(status) |
| 87 | +} |
| 88 | + |
| 89 | +/// Samples a normal (Gaussian) distribution with mean `loc` and standard |
| 90 | +/// deviation `scale`. |
| 91 | +/// |
| 92 | +/// `loc` and `scale` are `f32` because MLX's parameters are, regardless of the |
| 93 | +/// dtype chosen by `T`. Only floating-point dtypes can be generated, so |
| 94 | +/// `normal::<i32>` returns an error rather than truncating. |
| 95 | +pub fn normal<T: ArrayElement>( |
| 96 | + shape: &[i32], |
| 97 | + loc: f32, |
| 98 | + scale: f32, |
| 99 | + key: Option<&Key>, |
| 100 | + stream: &Stream, |
| 101 | +) -> Result<Array> { |
| 102 | + error::install(); |
| 103 | + let mut out = unsafe { sys::mlx_array_new() }; |
| 104 | + // SAFETY: `shape` describes a valid slice (or null/0); the key handle is |
| 105 | + // either a valid array or an absent one, both of which mlx accepts; the |
| 106 | + // result is written into `out`. |
| 107 | + let status = unsafe { |
| 108 | + sys::mlx_random_normal( |
| 109 | + &mut out, |
| 110 | + as_ffi_ptr(shape), |
| 111 | + shape.len(), |
| 112 | + T::DTYPE, |
| 113 | + loc, |
| 114 | + scale, |
| 115 | + Key::handle(key), |
| 116 | + stream.as_raw(), |
| 117 | + ) |
| 118 | + }; |
| 119 | + Array::from_op(out, status) |
| 120 | +} |
| 121 | + |
| 122 | +/// Samples a uniform distribution over `[low, high)`. |
| 123 | +/// |
| 124 | +/// Unlike [`normal`], MLX takes the bounds as arrays so they can broadcast; this |
| 125 | +/// wrapper accepts scalars of the element type and wraps them internally. |
| 126 | +pub fn uniform<T: ArrayElement>( |
| 127 | + shape: &[i32], |
| 128 | + low: T, |
| 129 | + high: T, |
| 130 | + key: Option<&Key>, |
| 131 | + stream: &Stream, |
| 132 | +) -> Result<Array> { |
| 133 | + error::install(); |
| 134 | + // Held in locals so they outlive the call below. |
| 135 | + let low = Array::from_scalar(low); |
| 136 | + let high = Array::from_scalar(high); |
| 137 | + let mut out = unsafe { sys::mlx_array_new() }; |
| 138 | + // SAFETY: the bound arrays outlive the call; `shape` describes a valid slice |
| 139 | + // (or null/0); the result is written into `out`. |
| 140 | + let status = unsafe { |
| 141 | + sys::mlx_random_uniform( |
| 142 | + &mut out, |
| 143 | + low.as_raw(), |
| 144 | + high.as_raw(), |
| 145 | + as_ffi_ptr(shape), |
| 146 | + shape.len(), |
| 147 | + T::DTYPE, |
| 148 | + Key::handle(key), |
| 149 | + stream.as_raw(), |
| 150 | + ) |
| 151 | + }; |
| 152 | + Array::from_op(out, status) |
| 153 | +} |
| 154 | + |
| 155 | +#[cfg(test)] |
| 156 | +mod tests { |
| 157 | + use super::*; |
| 158 | + |
| 159 | + #[test] |
| 160 | + fn same_key_gives_identical_draws() { |
| 161 | + let s = Stream::cpu(); |
| 162 | + let k1 = Key::new(42).unwrap(); |
| 163 | + let k2 = Key::new(42).unwrap(); |
| 164 | + let a = normal::<f32>(&[4], 0.0, 1.0, Some(&k1), &s).unwrap(); |
| 165 | + let b = normal::<f32>(&[4], 0.0, 1.0, Some(&k2), &s).unwrap(); |
| 166 | + assert_eq!(a.to_vec::<f32>(), b.to_vec::<f32>()); |
| 167 | + } |
| 168 | + |
| 169 | + #[test] |
| 170 | + fn different_seeds_give_different_draws() { |
| 171 | + let s = Stream::cpu(); |
| 172 | + let a = normal::<f32>(&[8], 0.0, 1.0, Some(&Key::new(1).unwrap()), &s).unwrap(); |
| 173 | + let b = normal::<f32>(&[8], 0.0, 1.0, Some(&Key::new(2).unwrap()), &s).unwrap(); |
| 174 | + assert_ne!(a.to_vec::<f32>(), b.to_vec::<f32>()); |
| 175 | + } |
| 176 | + |
| 177 | + #[test] |
| 178 | + fn normal_respects_shape_and_scale() { |
| 179 | + let s = Stream::cpu(); |
| 180 | + let key = Key::new(7).unwrap(); |
| 181 | + let a = normal::<f32>(&[2, 3], 0.0, 1.0, Some(&key), &s).unwrap(); |
| 182 | + assert_eq!(a.shape(), vec![2, 3]); |
| 183 | + assert_eq!(a.size(), 6); |
| 184 | + |
| 185 | + // A zero scale collapses the distribution onto `loc` exactly. |
| 186 | + let fixed = normal::<f32>(&[4], 5.0, 0.0, Some(&key), &s).unwrap(); |
| 187 | + assert_eq!(fixed.to_vec::<f32>(), vec![5.0; 4]); |
| 188 | + } |
| 189 | + |
| 190 | + #[test] |
| 191 | + fn uniform_stays_within_bounds() { |
| 192 | + let s = Stream::cpu(); |
| 193 | + let key = Key::new(3).unwrap(); |
| 194 | + let a = uniform::<f32>(&[64], -2.0, 2.0, Some(&key), &s).unwrap(); |
| 195 | + for v in a.to_vec::<f32>() { |
| 196 | + assert!((-2.0..2.0).contains(&v), "{v} out of bounds"); |
| 197 | + } |
| 198 | + } |
| 199 | + |
| 200 | + #[test] |
| 201 | + fn split_yields_independent_but_deterministic_keys() { |
| 202 | + let s = Stream::cpu(); |
| 203 | + let (a, b) = Key::new(11).unwrap().split(&s).unwrap(); |
| 204 | + let from_a = normal::<f32>(&[8], 0.0, 1.0, Some(&a), &s).unwrap(); |
| 205 | + let from_b = normal::<f32>(&[8], 0.0, 1.0, Some(&b), &s).unwrap(); |
| 206 | + // The halves are independent of each other... |
| 207 | + assert_ne!(from_a.to_vec::<f32>(), from_b.to_vec::<f32>()); |
| 208 | + |
| 209 | + // ...but splitting the same seed again reproduces the same pair. |
| 210 | + let (a2, _) = Key::new(11).unwrap().split(&s).unwrap(); |
| 211 | + let again = normal::<f32>(&[8], 0.0, 1.0, Some(&a2), &s).unwrap(); |
| 212 | + assert_eq!(from_a.to_vec::<f32>(), again.to_vec::<f32>()); |
| 213 | + } |
| 214 | + |
| 215 | + #[test] |
| 216 | + fn unkeyed_draws_use_global_state() { |
| 217 | + let s = Stream::cpu(); |
| 218 | + seed(99).unwrap(); |
| 219 | + // Deliberately only a structural assertion: the global state is shared |
| 220 | + // with every other test running in parallel, so the *values* here are |
| 221 | + // not reproducible. `Key` is the tool for that. |
| 222 | + let a = normal::<f32>(&[3], 0.0, 1.0, None, &s).unwrap(); |
| 223 | + assert_eq!(a.shape(), vec![3]); |
| 224 | + } |
| 225 | + |
| 226 | + #[test] |
| 227 | + fn integer_dtype_is_rejected() { |
| 228 | + let s = Stream::cpu(); |
| 229 | + let key = Key::new(1).unwrap(); |
| 230 | + // Normal draws are floating-point only; mlx reports this rather than |
| 231 | + // silently truncating. |
| 232 | + let err = normal::<i32>(&[4], 0.0, 1.0, Some(&key), &s).unwrap_err(); |
| 233 | + assert!( |
| 234 | + !err.message().is_empty(), |
| 235 | + "expected a message from mlx, got {err:?}" |
| 236 | + ); |
| 237 | + } |
| 238 | +} |
0 commit comments