mlx_rs/
random.rs

1//! Collection of functions related to random number generation
2
3use crate::ops::indexing::TryIndexOp;
4use crate::utils::guard::Guarded;
5use crate::utils::IntoOption;
6use crate::{error::Result, Array, ArrayElement, Stream};
7use mach_sys::mach_time;
8use mlx_internal_macros::generate_macro;
9use std::borrow::Cow;
10use std::cell::RefCell;
11
12thread_local! {
13    static THREAD_LOCAL_DEFAULT_STATE: RefCell<RandomState> = RefCell::new(RandomState::new().unwrap());
14    static THREAD_LOCAL_OVERRIDE_STATE: RefCell<Option<RandomState>> = const { RefCell::new(None) };
15}
16
17/// Random state for reproducible random number generation.
18///
19/// This struct holds the PRNG state and can be used with compiled functions
20/// to properly track random state across JIT compilation boundaries.
21///
22/// # Compilation Support
23///
24/// `RandomState` implements `Updatable`, making it compatible with
25/// `compile_with_state`. This is the Rust equivalent of Python's
26/// `@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)`.
27///
28/// # Example
29///
30/// ```rust,no_run
31/// use mlx_rs::random::RandomState;
32/// use mlx_rs::transforms::compile::compile_with_state;
33/// use mlx_rs::random::categorical;
34/// use mlx_rs::Array;
35///
36/// let mut state = RandomState::with_seed(42).unwrap();
37/// let logits = Array::zeros::<f32>(&[1, 10]).unwrap();
38/// let mut compiled = compile_with_state(
39///     |state: &mut RandomState, x: &Array| {
40///         let key = state.next_key()?;
41///         categorical(x, None, None, Some(&key))
42///     },
43///     None
44/// );
45/// let result = compiled(&mut state, &logits).unwrap();
46/// ```
47#[derive(Debug, Clone)]
48pub struct RandomState {
49    state: Array,
50}
51
52impl RandomState {
53    /// Create a new random state with a time-based seed.
54    pub fn new() -> Result<Self> {
55        let now = unsafe { mach_time::mach_approximate_time() };
56        Ok(Self { state: key(now)? })
57    }
58
59    /// Create a new random state from a specific seed.
60    ///
61    /// Use this for reproducible random number generation.
62    pub fn with_seed(seed: u64) -> Result<Self> {
63        Ok(Self { state: key(seed)? })
64    }
65
66    /// Create a random state from an existing key array.
67    ///
68    /// The key must be a valid PRNG key (typically created via `random::key()`).
69    pub fn from_key(key: Array) -> Self {
70        Self { state: key }
71    }
72
73    /// Get the next random key, advancing the state.
74    ///
75    /// This splits the current state into two keys: one becomes the new state,
76    /// and the other is returned for use in random operations.
77    pub fn next_key(&mut self) -> Result<Array> {
78        let next = split(&self.state, 2)?;
79        self.state = next.0;
80        Ok(next.1)
81    }
82
83    /// Internal method for backward compatibility.
84    fn next(&mut self) -> Result<Array> {
85        self.next_key()
86    }
87
88    /// Reseed the random state.
89    pub fn seed(&mut self, seed: u64) -> Result<()> {
90        self.state = key(seed)?;
91        Ok(())
92    }
93
94    /// Get a reference to the underlying state array.
95    ///
96    /// This is useful for inspection or manual state management.
97    pub fn as_array(&self) -> &Array {
98        &self.state
99    }
100
101    /// Get a mutable reference to the underlying state array.
102    ///
103    /// # Note
104    ///
105    /// Modifying the state array directly may break the PRNG invariants.
106    /// Prefer using `seed()` or `next_key()` instead.
107    pub fn as_array_mut(&mut self) -> &mut Array {
108        &mut self.state
109    }
110}
111
112impl Default for RandomState {
113    /// Creates a new `RandomState` with a time-based seed.
114    ///
115    /// # Panics
116    ///
117    /// Panics if the underlying PRNG key creation fails, which should not
118    /// occur under normal conditions.
119    fn default() -> Self {
120        Self::new().expect("Failed to create default RandomState")
121    }
122}
123
124impl crate::utils::Updatable for RandomState {
125    fn state_projection(
126        &mut self,
127    ) -> std::result::Result<crate::utils::StateProjection<'_>, crate::error::StateProjectionError>
128    {
129        let mut projection = crate::utils::StateProjection::new();
130        projection.required("key", &mut self.state)?;
131        Ok(projection)
132    }
133}
134
135/// Returns a key from the thread-local override state if it exists, otherwise
136/// returns `None`
137fn resolve_thread_local_override_key() -> Option<Result<Array>> {
138    THREAD_LOCAL_OVERRIDE_STATE.with_borrow_mut(|state| state.as_mut().map(|s| s.next()))
139}
140fn resolve_thread_local_default_key() -> Result<Array> {
141    THREAD_LOCAL_DEFAULT_STATE.with_borrow_mut(RandomState::next)
142}
143
144/// Use given key or generate a new one if `None`.
145fn resolve<'a>(key: impl Into<Option<&'a Array>>) -> Result<Cow<'a, Array>> {
146    key.into().map_or_else(
147        || {
148            resolve_thread_local_override_key()
149                .unwrap_or_else(resolve_thread_local_default_key)
150                .map(Cow::Owned)
151        },
152        |k| Ok(Cow::Borrowed(k)),
153    )
154}
155
156/// Use the given random state for the scope of `f`
157pub fn with_random_state<F, T>(state: RandomState, f: F) -> T
158where
159    F: FnOnce() -> T,
160{
161    let prev_state = THREAD_LOCAL_OVERRIDE_STATE.with_borrow_mut(|s| s.replace(state));
162
163    let result = f();
164
165    THREAD_LOCAL_OVERRIDE_STATE.with_borrow_mut(|s| {
166        *s = prev_state;
167    });
168
169    result
170}
171
172/// Seed the current thread's default random number generator.
173pub fn seed(seed: u64) -> Result<()> {
174    THREAD_LOCAL_DEFAULT_STATE.with_borrow_mut(|state| state.seed(seed))
175}
176
177/// Get a PRNG key from a seed.
178///
179/// Return a value that can be used as a PRNG key.  All ``random::*``
180/// functions take an optional key -- this will let you control the
181/// random number generation.
182pub fn key(seed: u64) -> Result<Array> {
183    Array::try_from_op(|res| unsafe { mlx_sys::mlx_random_key(res, seed) })
184}
185
186/// Split a PRNG key into two keys and return a tuple.
187pub fn split(key: impl AsRef<Array>, num: i32) -> Result<(Array, Array)> {
188    let stream = Stream::thread_local_or_default();
189    let keys = Array::try_from_op(|res| unsafe {
190        mlx_sys::mlx_random_split_num(res, key.as_ref().as_ptr(), num, stream.as_ref().as_ptr())
191    })?;
192
193    Ok((keys.try_index(0)?, keys.try_index(1)?))
194}
195
196/// Compatibility shim for [`split`].
197#[deprecated(
198    since = "0.26.0",
199    note = "use `with_stream` or `with_device` around `split`"
200)]
201pub fn split_device(
202    key: impl AsRef<Array>,
203    num: i32,
204    stream: impl AsRef<Stream>,
205) -> Result<(Array, Array)> {
206    crate::with_stream(stream.as_ref(), || split(key, num))
207}
208
209/// Generate uniformly distributed random numbers.
210/// The values are sampled uniformly in the half-open interval `[lower, upper)`.
211/// The lower and upper bound can be scalars or arrays and must be broadcastable to `shape`.
212///
213/// # Params
214///
215/// - `lower`: Lower bound of the distribution.
216/// - `upper`: Upper bound of the distribution.
217/// - `shape` (optional): Shape of the output. Default is `&[]`.
218/// - `key` (optional): A PRNG key.
219///
220/// ```rust
221/// let key = mlx_rs::random::key(0).unwrap();
222///
223/// // create an array of shape `[50]` type f32 values in the range [0, 10)
224/// let array = mlx_rs::random::uniform::<_, f32>(0, 10, &[50], &key);
225///
226/// // same, but in range [0.5, 1)
227/// let array = mlx_rs::random::uniform::<_, f32>(0.5f32, 1f32, &[50], &key);
228/// ```
229pub fn uniform<'a, E: Into<Array>, T: ArrayElement>(
230    lower: E,
231    upper: E,
232    shape: impl IntoOption<&'a [i32]>,
233    key: impl Into<Option<&'a Array>>,
234) -> Result<Array> {
235    let stream = Stream::thread_local_or_default();
236    let lb: Array = lower.into();
237    let ub: Array = upper.into();
238    let shape = shape.into_option().unwrap_or(&[]);
239    let key = resolve(key)?;
240
241    Array::try_from_op(|res| unsafe {
242        mlx_sys::mlx_random_uniform(
243            res,
244            lb.as_ptr(),
245            ub.as_ptr(),
246            shape.as_ptr(),
247            shape.len(),
248            T::DTYPE.into(),
249            key.as_ptr(),
250            stream.as_ref().as_ptr(),
251        )
252    })
253}
254
255/// Compatibility shim for [`uniform`].
256#[generate_macro(customize(forwarding_shim = true, root = "$crate::random"))]
257#[deprecated(
258    since = "0.26.0",
259    note = "use `with_stream` or `with_device` around `uniform`"
260)]
261pub fn uniform_device<'a, E: Into<Array>, T: ArrayElement>(
262    lower: E,
263    upper: E,
264    #[optional] shape: impl IntoOption<&'a [i32]>,
265    #[optional] key: impl Into<Option<&'a Array>>,
266    #[optional] stream: impl AsRef<Stream>,
267) -> Result<Array> {
268    crate::with_stream(stream.as_ref(), || {
269        uniform::<E, T>(lower, upper, shape, key)
270    })
271}
272
273/// Generate normally distributed random numbers.
274///
275/// Generate an array of random numbers using the optional shape. The result
276/// will be of the given `T`. `T` must be a floating point type.
277///
278/// # Params
279///
280///  - shape: shape of the output, if `None` a single value is returned
281///  - loc: mean of the distribution, default is `0.0`
282///  - scale: standard deviation of the distribution, default is `1.0`
283///  - key: PRNG key
284///
285/// # Example
286///
287/// ```rust
288/// let key = mlx_rs::random::key(0).unwrap();
289///
290/// // generate a single f32 with normal distribution
291/// let value = mlx_rs::random::normal::<f32>(None, None, None, &key).unwrap().item_exact::<f32>();
292///
293/// // generate an array of f32 with normal distribution in shape [10, 5]
294/// let array = mlx_rs::random::normal::<f32>(&[10, 5], None, None, &key);
295/// ```
296pub fn normal<'a, T: ArrayElement>(
297    shape: impl IntoOption<&'a [i32]>,
298    loc: impl Into<Option<f32>>,
299    scale: impl Into<Option<f32>>,
300    key: impl Into<Option<&'a Array>>,
301) -> Result<Array> {
302    let stream = Stream::thread_local_or_default();
303    let shape = shape.into_option().unwrap_or(&[]);
304    let key = resolve(key)?;
305
306    Array::try_from_op(|res| unsafe {
307        mlx_sys::mlx_random_normal(
308            res,
309            shape.as_ptr(),
310            shape.len(),
311            T::DTYPE.into(),
312            loc.into().unwrap_or(0.0),
313            scale.into().unwrap_or(1.0),
314            key.as_ptr(),
315            stream.as_ref().as_ptr(),
316        )
317    })
318}
319
320/// Compatibility shim for [`normal`].
321#[generate_macro(customize(forwarding_shim = true, root = "$crate::random"))]
322#[deprecated(
323    since = "0.26.0",
324    note = "use `with_stream` or `with_device` around `normal`"
325)]
326pub fn normal_device<'a, T: ArrayElement>(
327    #[optional] shape: impl IntoOption<&'a [i32]>,
328    #[optional] loc: impl Into<Option<f32>>,
329    #[optional] scale: impl Into<Option<f32>>,
330    #[optional] key: impl Into<Option<&'a Array>>,
331    #[optional] stream: impl AsRef<Stream>,
332) -> Result<Array> {
333    crate::with_stream(stream.as_ref(), || normal::<T>(shape, loc, scale, key))
334}
335
336/// Generate jointly-normal random samples given a mean and covariance.
337///
338/// The matrix `covariance` must be positive semi-definite. The behavior is
339/// undefined if it is not.  The only supported output type is f32.
340///
341/// # Params
342/// - `mean`: array of shape `[..., n]`, the mean of the distribution.
343/// - `covariance`: array  of shape `[..., n, n]`, the covariance matrix of the distribution. The batch shape `...` must be broadcast-compatible with that of `mean`.
344/// - `shape`: The output shape must be broadcast-compatible with `&mean.shape[..mean.shape.len()-1]` and `&covariance.shape[..covariance.shape.len()-2]`. If empty, the result shape is determined by broadcasting the batch shapes of `mean` and `covariance`.
345/// - `key`: PRNG key.
346// TODO: not supported on GPU yet
347pub fn multivariate_normal<'a, T: ArrayElement>(
348    mean: impl AsRef<Array>,
349    covariance: impl AsRef<Array>,
350    shape: impl IntoOption<&'a [i32]>,
351    key: impl Into<Option<&'a Array>>,
352) -> Result<Array> {
353    let stream = Stream::thread_local_or_cpu();
354    let shape = shape.into_option().unwrap_or(&[]);
355    let key = resolve(key)?;
356
357    Array::try_from_op(|res| unsafe {
358        mlx_sys::mlx_random_multivariate_normal(
359            res,
360            mean.as_ref().as_ptr(),
361            covariance.as_ref().as_ptr(),
362            shape.as_ptr(),
363            shape.len(),
364            T::DTYPE.into(),
365            key.as_ptr(),
366            stream.as_ref().as_ptr(),
367        )
368    })
369}
370
371/// Compatibility shim for [`multivariate_normal`].
372#[generate_macro(customize(root = "$crate::random"))]
373#[deprecated(
374    since = "0.26.0",
375    note = "use `with_stream` or `with_device` around `multivariate_normal`"
376)]
377pub fn multivariate_normal_device<'a, T: ArrayElement>(
378    mean: impl AsRef<Array>,
379    covariance: impl AsRef<Array>,
380    #[optional] shape: impl IntoOption<&'a [i32]>,
381    #[optional] key: impl Into<Option<&'a Array>>,
382    #[optional] stream: impl AsRef<Stream>,
383) -> Result<Array> {
384    crate::with_stream(stream.as_ref(), || {
385        multivariate_normal::<T>(mean, covariance, shape, key)
386    })
387}
388
389/// Generate random integers from the given interval (`lower:` and `upper:`).
390///
391/// The values are sampled with equal probability from the integers in
392/// half-open interval `[lb, ub)`. The lower and upper bound can be
393/// scalars or arrays and must be roadcastable to `shape`.
394///
395/// ```rust
396/// use mlx_rs::{array, random};
397///
398/// let key = random::key(0).unwrap();
399///
400/// // generate an array of Int values, one in the range [0, 20) and one in the range [10, 100)
401/// let array = random::randint::<_, i32>(array!([0, 20]), array!([10, 100]), None, &key);
402/// ```
403pub fn randint<'a, E: Into<Array>, T: ArrayElement>(
404    lower: E,
405    upper: E,
406    shape: impl IntoOption<&'a [i32]>,
407    key: impl Into<Option<&'a Array>>,
408) -> Result<Array> {
409    let stream = Stream::thread_local_or_default();
410    let lb: Array = lower.into();
411    let ub: Array = upper.into();
412    let shape = shape.into_option().unwrap_or(lb.shape());
413    let key = resolve(key)?;
414
415    Array::try_from_op(|res| unsafe {
416        mlx_sys::mlx_random_randint(
417            res,
418            lb.as_ptr(),
419            ub.as_ptr(),
420            shape.as_ptr(),
421            shape.len(),
422            T::DTYPE.into(),
423            key.as_ptr(),
424            stream.as_ref().as_ptr(),
425        )
426    })
427}
428
429/// Compatibility shim for [`randint`].
430#[generate_macro(customize(forwarding_shim = true, root = "$crate::random"))]
431#[deprecated(
432    since = "0.26.0",
433    note = "use `with_stream` or `with_device` around `randint`"
434)]
435pub fn randint_device<'a, E: Into<Array>, T: ArrayElement>(
436    lower: E,
437    upper: E,
438    #[optional] shape: impl IntoOption<&'a [i32]>,
439    #[optional] key: impl Into<Option<&'a Array>>,
440    #[optional] stream: impl AsRef<Stream>,
441) -> Result<Array> {
442    crate::with_stream(stream.as_ref(), || {
443        randint::<E, T>(lower, upper, shape, key)
444    })
445}
446
447/// Generate Bernoulli random values with a given `p` value.
448///
449/// The values are sampled from the bernoulli distribution with parameter
450/// `p`. The parameter `p` must have a floating point type and
451/// must be broadcastable to `shape`.
452///
453/// ```rust
454/// use mlx_rs::{array, Array, random};
455///
456/// let key = random::key(0).unwrap();
457///
458/// // generate a single random Bool with p = 0.8
459/// let p: Array = 0.8.into();
460/// let value = random::bernoulli(&p, None, &key);
461///
462/// // generate an array of shape [50, 2] of random Bool with p = 0.8
463/// let array = random::bernoulli(&p, &[50, 2], &key);
464///
465/// // generate an array of [3] Bool with the given p values
466/// let array = random::bernoulli(&array!([0.1, 0.5, 0.8]), None, &key);
467/// ```
468pub fn bernoulli<'a>(
469    p: impl Into<Option<&'a Array>>,
470    shape: impl IntoOption<&'a [i32]>,
471    key: impl Into<Option<&'a Array>>,
472) -> Result<Array> {
473    let stream = Stream::thread_local_or_default();
474    let default_array = Array::from_f32(0.5);
475    let p = p.into().unwrap_or(&default_array);
476
477    let shape = shape.into_option().unwrap_or(p.shape());
478    let key = resolve(key)?;
479
480    Array::try_from_op(|res| unsafe {
481        mlx_sys::mlx_random_bernoulli(
482            res,
483            p.as_ptr(),
484            shape.as_ptr(),
485            shape.len(),
486            key.as_ptr(),
487            stream.as_ref().as_ptr(),
488        )
489    })
490}
491
492/// Compatibility shim for [`bernoulli`].
493#[generate_macro(customize(forwarding_shim = true, root = "$crate::random"))]
494#[deprecated(
495    since = "0.26.0",
496    note = "use `with_stream` or `with_device` around `bernoulli`"
497)]
498pub fn bernoulli_device<'a>(
499    #[optional] p: impl Into<Option<&'a Array>>,
500    #[optional] shape: impl IntoOption<&'a [i32]>,
501    #[optional] key: impl Into<Option<&'a Array>>,
502    #[optional] stream: impl AsRef<Stream>,
503) -> Result<Array> {
504    crate::with_stream(stream.as_ref(), || bernoulli(p, shape, key))
505}
506
507/// Generate values from a truncated normal distribution between `low` and `high`.
508///
509/// The values are sampled from the truncated normal distribution
510/// on the domain `(lower, upper)`. The bounds `lower` and `upper`
511/// can be scalars or arrays and must be broadcastable to `shape`.
512///
513/// ```rust
514/// use mlx_rs::{array, random};
515///
516/// let key = random::key(0).unwrap();
517///
518/// // generate an array of two Float values, one in the range 0 ..< 10
519/// // and one in the range 10 ..< 100
520/// let value = random::truncated_normal::<_, f32>(array!([0, 10]), array!([10, 100]), None, &key);
521/// ```
522pub fn truncated_normal<'a, E: Into<Array>, T: ArrayElement>(
523    lower: E,
524    upper: E,
525    shape: impl IntoOption<&'a [i32]>,
526    key: impl Into<Option<&'a Array>>,
527) -> Result<Array> {
528    let stream = Stream::thread_local_or_default();
529    let lb: Array = lower.into();
530    let ub: Array = upper.into();
531    let shape = shape.into_option().unwrap_or(lb.shape());
532    let key = resolve(key)?;
533
534    Array::try_from_op(|res| unsafe {
535        mlx_sys::mlx_random_truncated_normal(
536            res,
537            lb.as_ptr(),
538            ub.as_ptr(),
539            shape.as_ptr(),
540            shape.len(),
541            T::DTYPE.into(),
542            key.as_ptr(),
543            stream.as_ref().as_ptr(),
544        )
545    })
546}
547
548/// Compatibility shim for [`truncated_normal`].
549#[generate_macro(customize(forwarding_shim = true, root = "$crate::random"))]
550#[deprecated(
551    since = "0.26.0",
552    note = "use `with_stream` or `with_device` around `truncated_normal`"
553)]
554pub fn truncated_normal_device<'a, E: Into<Array>, T: ArrayElement>(
555    lower: E,
556    upper: E,
557    #[optional] shape: impl IntoOption<&'a [i32]>,
558    #[optional] key: impl Into<Option<&'a Array>>,
559    #[optional] stream: impl AsRef<Stream>,
560) -> Result<Array> {
561    crate::with_stream(stream.as_ref(), || {
562        truncated_normal::<E, T>(lower, upper, shape, key)
563    })
564}
565
566/// Sample from the standard Gumbel distribution.
567///
568/// The values are sampled from a standard Gumbel distribution
569/// which CDF `exp(-exp(-x))`.
570///
571/// ```rust
572/// let key = mlx_rs::random::key(0).unwrap();
573///
574/// // generate a single Float with Gumbel distribution
575/// let value = mlx_rs::random::gumbel::<f32>(None, &key).unwrap().item_exact::<f32>();
576///
577/// // generate an array of Float with Gumbel distribution in shape [10, 5]
578/// let array = mlx_rs::random::gumbel::<f32>(&[10, 5], &key);
579/// ```
580pub fn gumbel<'a, T: ArrayElement>(
581    shape: impl IntoOption<&'a [i32]>,
582    key: impl Into<Option<&'a Array>>,
583) -> Result<Array> {
584    let stream = Stream::thread_local_or_default();
585    let shape = shape.into_option().unwrap_or(&[]);
586    let key = resolve(key)?;
587
588    Array::try_from_op(|res| unsafe {
589        mlx_sys::mlx_random_gumbel(
590            res,
591            shape.as_ptr(),
592            shape.len(),
593            T::DTYPE.into(),
594            key.as_ptr(),
595            stream.as_ref().as_ptr(),
596        )
597    })
598}
599
600/// Compatibility shim for [`gumbel`].
601#[generate_macro(customize(forwarding_shim = true, root = "$crate::random"))]
602#[deprecated(
603    since = "0.26.0",
604    note = "use `with_stream` or `with_device` around `gumbel`"
605)]
606pub fn gumbel_device<'a, T: ArrayElement>(
607    #[optional] shape: impl IntoOption<&'a [i32]>,
608    #[optional] key: impl Into<Option<&'a Array>>,
609    #[optional] stream: impl AsRef<Stream>,
610) -> Result<Array> {
611    crate::with_stream(stream.as_ref(), || gumbel::<T>(shape, key))
612}
613
614/// Shape or count for the categorical distribution.
615#[derive(Debug, Clone, Copy)]
616pub enum ShapeOrCount<'a> {
617    /// Shape
618    Shape(&'a [i32]),
619
620    /// Count
621    Count(i32),
622}
623
624/// Sample from a categorical distribution.
625///
626/// The values are sampled from the categorical distribution specified by
627/// the unnormalized values in `logits`.   If the `shape` is not specified
628/// the result shape will be the same shape as `logits` with the `axis`
629/// dimension removed.
630///
631/// /// # Params
632/// # Params
633///
634/// - `logits`: The *unnormalized* categorical distribution(s).
635/// - `axis`(optional): The axis which specifies the distribution. Default is `-1`.
636/// - `shape_or_count`(optional):
637/// - - `Shape`: The shape of the output. This must be broadcast compatible with `logits.shape` with the `axis` dimension removed.
638/// - - `Count`: The number of samples to draw from each of the categorical distributions in `logits`. The output will have the number of samples in the last dimension.
639/// - `key` (optional): A PRNG key.
640///
641/// # Example
642///
643/// ```rust
644/// let key = mlx_rs::random::key(0).unwrap();
645///
646/// let logits = mlx_rs::Array::zeros::<u32>(&[5, 20]).unwrap();
647///
648/// // produces Array of u32 shape &[5]
649/// let result = mlx_rs::random::categorical(&logits, None, None, &key);
650/// ```
651pub fn categorical<'a>(
652    logits: impl AsRef<Array>,
653    axis: impl Into<Option<i32>>,
654    shape_or_count: impl Into<Option<ShapeOrCount<'a>>>,
655    key: impl Into<Option<&'a Array>>,
656) -> Result<Array> {
657    let stream = Stream::thread_local_or_default();
658    let axis = axis.into().unwrap_or(-1);
659    let key = resolve(key)?;
660
661    match shape_or_count.into() {
662        Some(ShapeOrCount::Shape(shape)) => Array::try_from_op(|res| unsafe {
663            mlx_sys::mlx_random_categorical_shape(
664                res,
665                logits.as_ref().as_ptr(),
666                axis,
667                shape.as_ptr(),
668                shape.len(),
669                key.as_ptr(),
670                stream.as_ref().as_ptr(),
671            )
672        }),
673        Some(ShapeOrCount::Count(num_samples)) => Array::try_from_op(|res| unsafe {
674            mlx_sys::mlx_random_categorical_num_samples(
675                res,
676                logits.as_ref().as_ptr(),
677                axis,
678                num_samples,
679                key.as_ptr(),
680                stream.as_ref().as_ptr(),
681            )
682        }),
683        None => Array::try_from_op(|res| unsafe {
684            mlx_sys::mlx_random_categorical(
685                res,
686                logits.as_ref().as_ptr(),
687                axis,
688                key.as_ptr(),
689                stream.as_ref().as_ptr(),
690            )
691        }),
692    }
693}
694
695/// Compatibility shim for [`categorical`].
696#[generate_macro(customize(forwarding_shim = true, root = "$crate::random"))]
697#[deprecated(
698    since = "0.26.0",
699    note = "use `with_stream` or `with_device` around `categorical`"
700)]
701pub fn categorical_device<'a>(
702    logits: impl AsRef<Array>,
703    #[optional] axis: impl Into<Option<i32>>,
704    #[optional] shape_or_count: impl Into<Option<ShapeOrCount<'a>>>,
705    #[optional] key: impl Into<Option<&'a Array>>,
706    #[optional] stream: impl AsRef<Stream>,
707) -> Result<Array> {
708    crate::with_stream(stream.as_ref(), || {
709        categorical(logits, axis, shape_or_count, key)
710    })
711}
712
713#[cfg(test)]
714mod tests {
715    use super::*;
716    use crate::{
717        array,
718        test_utils::{assert_array_eq, assert_array_eq_with_context, tolerances},
719    };
720    use float_eq::{assert_float_eq, float_eq};
721
722    #[test]
723    fn test_default_rng() {
724        seed(3).unwrap();
725        let a = uniform::<_, f32>(0, 1, None, None).unwrap();
726        let b = uniform::<_, f32>(0, 1, None, None).unwrap();
727
728        seed(3).unwrap();
729        let x = uniform::<_, f32>(0, 1, None, None).unwrap();
730        let y = uniform::<_, f32>(0, 1, None, None).unwrap();
731
732        assert_array_eq(a, x, tolerances::EXACT.rtol, tolerances::EXACT.atol);
733        assert_array_eq(b, y, tolerances::EXACT.rtol, tolerances::EXACT.atol);
734    }
735
736    #[test]
737    fn sequential_threads_use_own_default_rng_stream() {
738        crate::Device::set_default(&crate::Device::gpu());
739
740        for _ in 0..2 {
741            std::thread::spawn(|| {
742                uniform::<_, f32>(0.0, 1.0, &[8], None)
743                    .unwrap()
744                    .eval()
745                    .unwrap();
746            })
747            .join()
748            .unwrap();
749        }
750    }
751
752    #[test]
753    fn test_key() {
754        let k1 = key(0).unwrap();
755        let k2 = key(0).unwrap();
756        assert_array_eq(&k1, k2, tolerances::EXACT.rtol, tolerances::EXACT.atol);
757
758        let k2 = key(1).unwrap();
759        assert!(!k1.eq_exact(&k2).unwrap());
760    }
761
762    #[test]
763    fn test_split() {
764        let key = key(0).unwrap();
765
766        let (k1, k2) = split(&key, 2).unwrap();
767        assert!(!k1.eq_exact(&k2).unwrap());
768
769        let (r1, r2) = split(&key, 2).unwrap();
770        assert_array_eq(r1, k1, tolerances::EXACT.rtol, tolerances::EXACT.atol);
771        assert_array_eq(r2, k2, tolerances::EXACT.rtol, tolerances::EXACT.atol);
772    }
773
774    #[test]
775    fn test_uniform_no_seed() {
776        let value = uniform::<_, f32>(0, 10, &[3], None).unwrap();
777        assert_eq!(value.shape(), &[3]);
778    }
779
780    #[test]
781    fn test_uniform_single() {
782        let key = key(0).unwrap();
783        let value = uniform::<_, f32>(0, 10, None, Some(&key)).unwrap();
784        float_eq!(value.item_exact::<f32>(), 4.18, abs <= 0.01);
785    }
786
787    #[test]
788    fn test_uniform_multiple() {
789        let key = key(0).unwrap();
790        let value = uniform::<_, f32>(0, 10, &[3], Some(&key)).unwrap();
791        let expected = Array::from_slice(&[9.65, 3.14, 6.33], &[3]);
792
793        assert_array_eq(
794            value,
795            expected,
796            tolerances::ROUNDED_TWO_DECIMALS.rtol,
797            tolerances::ROUNDED_TWO_DECIMALS.atol,
798        );
799    }
800
801    #[test]
802    fn test_uniform_multiple_array() {
803        let key = key(0).unwrap();
804        let value = uniform::<_, f32>(&[0, 10], &[10, 100], &[2], Some(&key)).unwrap();
805        let expected = Array::from_slice(&[2.16, 82.37], &[2]);
806
807        assert_array_eq(
808            value,
809            expected,
810            tolerances::ROUNDED_TWO_DECIMALS.rtol,
811            tolerances::ROUNDED_TWO_DECIMALS.atol,
812        );
813    }
814
815    #[test]
816    fn test_uniform_non_float() {
817        let key = key(0).unwrap();
818        let value = uniform::<_, i32>(&[0, 10], &[10, 100], &[2], Some(&key));
819        assert!(value.is_err());
820    }
821
822    #[test]
823    fn test_normal() {
824        let key = key(0).unwrap();
825        let value = normal::<f32>(None, None, None, &key).unwrap();
826        float_eq!(value.item_exact::<f32>(), -0.20, abs <= 0.01);
827    }
828
829    #[test]
830    fn test_normal_non_float() {
831        let key = key(0).unwrap();
832        let value = normal::<i32>(None, None, None, &key);
833        assert!(value.is_err());
834    }
835
836    #[test]
837    fn test_multivariate_normal() {
838        let key = key(0).unwrap();
839        let mean = Array::from_slice(&[0.0, 0.0], &[2]);
840        let covariance = Array::from_slice(&[1.0, 0.0, 0.0, 1.0], &[2, 2]);
841
842        let a = multivariate_normal::<f32>(&mean, &covariance, &[3], &key).unwrap();
843        assert!(a.shape() == [3, 2]);
844    }
845
846    #[test]
847    fn test_randint_single() {
848        let key = key(0).unwrap();
849        let value = randint::<_, i32>(0, 100, None, Some(&key)).unwrap();
850        assert_eq!(value.item_exact::<i32>(), 41);
851    }
852
853    #[test]
854    fn test_randint_multiple() {
855        let key = key(0).unwrap();
856        let value =
857            randint::<_, i32>(array!([0, 10]), array!([10, 100]), None, Some(&key)).unwrap();
858        let expected = Array::from_slice(&[2, 82], &[2]);
859
860        assert_array_eq(
861            value,
862            expected,
863            tolerances::EXACT.rtol,
864            tolerances::EXACT.atol,
865        );
866    }
867
868    #[test]
869    fn test_randint_non_int() {
870        let key = key(0).unwrap();
871        let value = randint::<_, f32>(array!([0, 10]), array!([10, 100]), None, Some(&key));
872        assert!(value.is_err());
873    }
874
875    #[test]
876    fn test_bernoulli_single() {
877        let key = key(0).unwrap();
878        let value = bernoulli(None, None, &key).unwrap();
879        assert!(value.item_exact::<bool>());
880    }
881
882    #[test]
883    fn test_bernoulli_multiple() {
884        let key = key(0).unwrap();
885        let value = bernoulli(None, &[4], &key).unwrap();
886        let expected = Array::from_slice(&[false, true, false, true], &[4]);
887
888        assert_array_eq(
889            value,
890            expected,
891            tolerances::EXACT.rtol,
892            tolerances::EXACT.atol,
893        );
894    }
895
896    #[test]
897    fn test_bernoulli_p() {
898        let key = key(0).unwrap();
899        let p: Array = 0.8.into();
900        let value = bernoulli(&p, &[4], &key).unwrap();
901        let expected = Array::from_slice(&[false, true, true, true], &[4]);
902
903        assert_array_eq(
904            value,
905            expected,
906            tolerances::EXACT.rtol,
907            tolerances::EXACT.atol,
908        );
909    }
910
911    #[test]
912    fn test_bernoulli_p_array() {
913        let key = key(0).unwrap();
914        let value = bernoulli(&array!([0.1, 0.5, 0.8]), None, &key).unwrap();
915        let expected = Array::from_slice(&[false, true, true], &[3]);
916
917        assert_array_eq(
918            value,
919            expected,
920            tolerances::EXACT.rtol,
921            tolerances::EXACT.atol,
922        );
923    }
924
925    #[test]
926    fn test_truncated_normal_single() {
927        let key = key(0).unwrap();
928        let value = truncated_normal::<_, f32>(0, 10, None, &key).unwrap();
929        assert_array_eq(
930            value,
931            Array::from_f32(0.55),
932            tolerances::ROUNDED_TWO_DECIMALS.rtol,
933            tolerances::ROUNDED_TWO_DECIMALS.atol,
934        );
935    }
936
937    #[test]
938    fn test_truncated_normal_multiple() {
939        let key = key(0).unwrap();
940        let value = truncated_normal::<_, f32>(0.0, 0.5, &[3], &key).unwrap();
941        let expected = Array::from_slice(&[0.48, 0.15, 0.30], &[3]);
942
943        assert_array_eq(
944            value,
945            expected,
946            tolerances::ROUNDED_TWO_DECIMALS.rtol,
947            tolerances::ROUNDED_TWO_DECIMALS.atol,
948        );
949    }
950
951    #[test]
952    fn test_truncated_normal_multiple_array() {
953        let key = key(0).unwrap();
954        let value =
955            truncated_normal::<_, f32>(array!([0.0, 0.5]), array!([0.5, 1.0]), None, &key).unwrap();
956        let expected = Array::from_slice(&[0.10, 0.88], &[2]);
957
958        assert_array_eq(
959            value,
960            expected,
961            tolerances::ROUNDED_TWO_DECIMALS.rtol,
962            tolerances::ROUNDED_TWO_DECIMALS.atol,
963        );
964    }
965
966    #[test]
967    fn test_gumbel() {
968        let key = key(0).unwrap();
969        let value = gumbel::<f32>(None, &key).unwrap();
970        assert_array_eq(
971            value,
972            Array::from_f32(0.13),
973            tolerances::ROUNDED_TWO_DECIMALS.rtol,
974            tolerances::ROUNDED_TWO_DECIMALS.atol,
975        );
976    }
977
978    #[test]
979    fn test_logits() {
980        let key = key(0).unwrap();
981        let logits = Array::zeros::<u32>(&[5, 20]).unwrap();
982        let result = categorical(&logits, None, None, &key).unwrap();
983
984        assert_eq!(result.shape(), [5]);
985
986        let expected = Array::from_slice(&[1_u32, 1, 17, 17, 17], &[5]);
987        assert_array_eq_with_context(
988            result,
989            expected,
990            tolerances::EXACT.rtol,
991            tolerances::EXACT.atol,
992            "categorical default sample values",
993        );
994    }
995
996    #[test]
997    fn test_logits_count() {
998        let key = key(0).unwrap();
999        let logits = Array::zeros::<u32>(&[5, 20]).unwrap();
1000        let result = categorical(&logits, None, ShapeOrCount::Count(2), &key).unwrap();
1001
1002        assert_eq!(result.shape(), [5, 2]);
1003
1004        let expected = Array::from_slice(&[16_u32, 3, 14, 10, 17, 7, 6, 8, 12, 8], &[5, 2]);
1005        assert_array_eq_with_context(
1006            result,
1007            expected,
1008            tolerances::EXACT.rtol,
1009            tolerances::EXACT.atol,
1010            "categorical counted sample values",
1011        );
1012    }
1013
1014    #[test]
1015    fn test_random_state_new() {
1016        let state = RandomState::new().unwrap();
1017        assert_eq!(state.as_array().shape(), &[2]);
1018    }
1019
1020    #[test]
1021    fn test_random_state_with_seed_deterministic() {
1022        let s1 = RandomState::with_seed(42).unwrap();
1023        let s2 = RandomState::with_seed(42).unwrap();
1024        assert_array_eq(
1025            s1.as_array(),
1026            s2.as_array(),
1027            tolerances::EXACT.rtol,
1028            tolerances::EXACT.atol,
1029        );
1030    }
1031
1032    #[test]
1033    fn test_random_state_next_key_advances() {
1034        let mut state = RandomState::with_seed(0).unwrap();
1035        let k1 = state.next_key().unwrap();
1036        let k2 = state.next_key().unwrap();
1037        assert!(!k1.eq_exact(&k2).unwrap());
1038    }
1039
1040    #[test]
1041    fn test_random_state_from_key_roundtrip() {
1042        let original = RandomState::with_seed(99).unwrap();
1043        let arr = original.as_array().clone();
1044        let restored = RandomState::from_key(arr);
1045        assert_array_eq(
1046            original.as_array(),
1047            restored.as_array(),
1048            tolerances::EXACT.rtol,
1049            tolerances::EXACT.atol,
1050        );
1051    }
1052
1053    #[test]
1054    fn test_random_state_updatable() {
1055        use crate::utils::Updatable;
1056        let mut state = RandomState::with_seed(0).unwrap();
1057        let projection = state.state_projection().unwrap();
1058        assert_eq!(projection.len(), 1);
1059        assert_eq!(projection.values().count(), 1);
1060    }
1061
1062    #[test]
1063    fn test_random_state_default() {
1064        let state = RandomState::default();
1065        assert_eq!(state.as_array().shape(), &[2]);
1066    }
1067
1068    #[test]
1069    fn test_random_seed_same() {
1070        // Same random seed should produce the same results
1071        let seed = 23;
1072        let mut results = Vec::new();
1073        let f = || {
1074            let sum = uniform::<_, f32>(0.0, 1.0, &[10, 10], None)?.sum(None)?;
1075            Ok::<_, crate::error::Exception>(sum.item_exact::<f32>())
1076        };
1077        for _ in 0..10 {
1078            let mut state = RandomState::new().unwrap();
1079            state.seed(seed).unwrap();
1080            let result = with_random_state(state, f).unwrap();
1081            results.push(result);
1082        }
1083
1084        // Check that all results are the same within a small tolerance
1085        let first = results[0];
1086        for result in &results[1..] {
1087            assert_float_eq!(
1088                first,
1089                *result,
1090                abs <= 0.01,
1091                "Results should be equal for the same seed"
1092            );
1093        }
1094    }
1095}
mlx_rs/random.rs

mlx_rs/
random.rs
