Skip to content

Commit 75fbf9e

Browse files
authored
support RNG (#9)
Signed-off-by: kerthcet <kerthcet@gmail.com>
1 parent ed0d707 commit 75fbf9e

4 files changed

Lines changed: 271 additions & 15 deletions

File tree

‎crates/mlxcore/src/array.rs‎

Lines changed: 2 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use mlxcore_sys as sys;
77

88
use crate::dtype::ArrayElement;
99
use crate::error::{self, Result};
10+
use crate::ffi::as_ffi_ptr;
1011
use crate::stream::Stream;
1112

1213
/// An N-dimensional MLX array.
@@ -16,20 +17,6 @@ pub struct Array {
1617
handle: sys::mlx_array,
1718
}
1819

19-
/// Returns a pointer to `slice`'s data, or null when it is empty.
20-
///
21-
/// For an empty slice `as_ptr()` is non-null but dangling. The mlx-c functions
22-
/// we call take a `(ptr, len)` pair and build a container from it, so they never
23-
/// dereference the pointer when `len == 0` — but passing an explicit null keeps
24-
/// a dangling pointer from crossing the FFI boundary at all.
25-
fn as_ffi_ptr<T>(slice: &[T]) -> *const T {
26-
if slice.is_empty() {
27-
std::ptr::null()
28-
} else {
29-
slice.as_ptr()
30-
}
31-
}
32-
3320
impl Array {
3421
/// Creates an `Array` from a raw handle, taking ownership of it.
3522
///
@@ -634,7 +621,7 @@ impl Array {
634621
///
635622
/// On failure, frees the (unused) `out` handle and returns the captured
636623
/// MLX error message.
637-
fn from_op(out: sys::mlx_array, status: i32) -> Result<Array> {
624+
pub(crate) fn from_op(out: sys::mlx_array, status: i32) -> Result<Array> {
638625
match error::check(status) {
639626
Ok(()) => Ok(unsafe { Self::from_raw(out) }),
640627
Err(e) => {

‎crates/mlxcore/src/ffi.rs‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
//! Small helpers shared by the modules that call into `mlxcore-sys`.
2+
3+
use mlxcore_sys as sys;
4+
5+
/// Returns a pointer to `slice`'s data, or null when it is empty.
6+
///
7+
/// For an empty slice `as_ptr()` is non-null but dangling. The mlx-c functions
8+
/// we call take a `(ptr, len)` pair and build a container from it, so they never
9+
/// dereference the pointer when `len == 0` — but passing an explicit null keeps
10+
/// a dangling pointer from crossing the FFI boundary at all.
11+
pub(crate) fn as_ffi_ptr<T>(slice: &[T]) -> *const T {
12+
if slice.is_empty() {
13+
std::ptr::null()
14+
} else {
15+
slice.as_ptr()
16+
}
17+
}
18+
19+
/// An absent `mlx_array`, for the parameters mlx-c documents as `may be null`.
20+
///
21+
/// The C shims test the handle's `ctx` pointer and translate a null one into
22+
/// `std::nullopt`, so this is how an optional array argument is omitted. It owns
23+
/// nothing and must not be freed.
24+
pub(crate) fn absent_array() -> sys::mlx_array {
25+
sys::mlx_array {
26+
ctx: std::ptr::null_mut(),
27+
}
28+
}

‎crates/mlxcore/src/lib.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,11 @@
66
mod array;
77
mod dtype;
88
mod error;
9+
mod ffi;
910
mod stream;
1011

12+
pub mod random;
13+
1114
pub use array::Array;
1215
pub use dtype::ArrayElement;
1316
pub use error::{Error, Result};

‎crates/mlxcore/src/random.rs‎

Lines changed: 238 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,238 @@
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

Comments
 (0)