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    #[cfg(feature = "metal")]
738    fn sequential_threads_use_own_default_rng_stream() {
739        crate::Device::set_default(&crate::Device::gpu());
740
741        for _ in 0..2 {
742            std::thread::spawn(|| {
743                uniform::<_, f32>(0.0, 1.0, &[8], None)
744                    .unwrap()
745                    .eval()
746                    .unwrap();
747            })
748            .join()
749            .unwrap();
750        }
751    }
752
753    #[test]
754    fn test_key() {
755        let k1 = key(0).unwrap();
756        let k2 = key(0).unwrap();
757        assert_array_eq(&k1, k2, tolerances::EXACT.rtol, tolerances::EXACT.atol);
758
759        let k2 = key(1).unwrap();
760        assert!(!k1.eq_exact(&k2).unwrap());
761    }
762
763    #[test]
764    fn test_split() {
765        let key = key(0).unwrap();
766
767        let (k1, k2) = split(&key, 2).unwrap();
768        assert!(!k1.eq_exact(&k2).unwrap());
769
770        let (r1, r2) = split(&key, 2).unwrap();
771        assert_array_eq(r1, k1, tolerances::EXACT.rtol, tolerances::EXACT.atol);
772        assert_array_eq(r2, k2, tolerances::EXACT.rtol, tolerances::EXACT.atol);
773    }
774
775    #[test]
776    fn test_uniform_no_seed() {
777        let value = uniform::<_, f32>(0, 10, &[3], None).unwrap();
778        assert_eq!(value.shape(), &[3]);
779    }
780
781    #[test]
782    fn test_uniform_single() {
783        let key = key(0).unwrap();
784        let value = uniform::<_, f32>(0, 10, None, Some(&key)).unwrap();
785        float_eq!(value.item_exact::<f32>(), 4.18, abs <= 0.01);
786    }
787
788    #[test]
789    fn test_uniform_multiple() {
790        let key = key(0).unwrap();
791        let value = uniform::<_, f32>(0, 10, &[3], Some(&key)).unwrap();
792        let expected = Array::from_slice(&[9.65, 3.14, 6.33], &[3]);
793
794        assert_array_eq(
795            value,
796            expected,
797            tolerances::ROUNDED_TWO_DECIMALS.rtol,
798            tolerances::ROUNDED_TWO_DECIMALS.atol,
799        );
800    }
801
802    #[test]
803    fn test_uniform_multiple_array() {
804        let key = key(0).unwrap();
805        let value = uniform::<_, f32>(&[0, 10], &[10, 100], &[2], Some(&key)).unwrap();
806        let expected = Array::from_slice(&[2.16, 82.37], &[2]);
807
808        assert_array_eq(
809            value,
810            expected,
811            tolerances::ROUNDED_TWO_DECIMALS.rtol,
812            tolerances::ROUNDED_TWO_DECIMALS.atol,
813        );
814    }
815
816    #[test]
817    fn test_uniform_non_float() {
818        let key = key(0).unwrap();
819        let value = uniform::<_, i32>(&[0, 10], &[10, 100], &[2], Some(&key));
820        assert!(value.is_err());
821    }
822
823    #[test]
824    fn test_normal() {
825        let key = key(0).unwrap();
826        let value = normal::<f32>(None, None, None, &key).unwrap();
827        float_eq!(value.item_exact::<f32>(), -0.20, abs <= 0.01);
828    }
829
830    #[test]
831    fn test_normal_non_float() {
832        let key = key(0).unwrap();
833        let value = normal::<i32>(None, None, None, &key);
834        assert!(value.is_err());
835    }
836
837    #[test]
838    fn test_multivariate_normal() {
839        let key = key(0).unwrap();
840        let mean = Array::from_slice(&[0.0, 0.0], &[2]);
841        let covariance = Array::from_slice(&[1.0, 0.0, 0.0, 1.0], &[2, 2]);
842
843        let a = multivariate_normal::<f32>(&mean, &covariance, &[3], &key).unwrap();
844        assert!(a.shape() == [3, 2]);
845    }
846
847    #[test]
848    fn test_randint_single() {
849        let key = key(0).unwrap();
850        let value = randint::<_, i32>(0, 100, None, Some(&key)).unwrap();
851        assert_eq!(value.item_exact::<i32>(), 41);
852    }
853
854    #[test]
855    fn test_randint_multiple() {
856        let key = key(0).unwrap();
857        let value =
858            randint::<_, i32>(array!([0, 10]), array!([10, 100]), None, Some(&key)).unwrap();
859        let expected = Array::from_slice(&[2, 82], &[2]);
860
861        assert_array_eq(
862            value,
863            expected,
864            tolerances::EXACT.rtol,
865            tolerances::EXACT.atol,
866        );
867    }
868
869    #[test]
870    fn test_randint_non_int() {
871        let key = key(0).unwrap();
872        let value = randint::<_, f32>(array!([0, 10]), array!([10, 100]), None, Some(&key));
873        assert!(value.is_err());
874    }
875
876    #[test]
877    fn test_bernoulli_single() {
878        let key = key(0).unwrap();
879        let value = bernoulli(None, None, &key).unwrap();
880        assert!(value.item_exact::<bool>());
881    }
882
883    #[test]
884    fn test_bernoulli_multiple() {
885        let key = key(0).unwrap();
886        let value = bernoulli(None, &[4], &key).unwrap();
887        let expected = Array::from_slice(&[false, true, false, true], &[4]);
888
889        assert_array_eq(
890            value,
891            expected,
892            tolerances::EXACT.rtol,
893            tolerances::EXACT.atol,
894        );
895    }
896
897    #[test]
898    fn test_bernoulli_p() {
899        let key = key(0).unwrap();
900        let p: Array = 0.8.into();
901        let value = bernoulli(&p, &[4], &key).unwrap();
902        let expected = Array::from_slice(&[false, true, true, true], &[4]);
903
904        assert_array_eq(
905            value,
906            expected,
907            tolerances::EXACT.rtol,
908            tolerances::EXACT.atol,
909        );
910    }
911
912    #[test]
913    fn test_bernoulli_p_array() {
914        let key = key(0).unwrap();
915        let value = bernoulli(&array!([0.1, 0.5, 0.8]), None, &key).unwrap();
916        let expected = Array::from_slice(&[false, true, true], &[3]);
917
918        assert_array_eq(
919            value,
920            expected,
921            tolerances::EXACT.rtol,
922            tolerances::EXACT.atol,
923        );
924    }
925
926    #[test]
927    fn test_truncated_normal_single() {
928        let key = key(0).unwrap();
929        let value = truncated_normal::<_, f32>(0, 10, None, &key).unwrap();
930        assert_array_eq(
931            value,
932            Array::from_f32(0.55),
933            tolerances::ROUNDED_TWO_DECIMALS.rtol,
934            tolerances::ROUNDED_TWO_DECIMALS.atol,
935        );
936    }
937
938    #[test]
939    fn test_truncated_normal_multiple() {
940        let key = key(0).unwrap();
941        let value = truncated_normal::<_, f32>(0.0, 0.5, &[3], &key).unwrap();
942        let expected = Array::from_slice(&[0.48, 0.15, 0.30], &[3]);
943
944        assert_array_eq(
945            value,
946            expected,
947            tolerances::ROUNDED_TWO_DECIMALS.rtol,
948            tolerances::ROUNDED_TWO_DECIMALS.atol,
949        );
950    }
951
952    #[test]
953    fn test_truncated_normal_multiple_array() {
954        let key = key(0).unwrap();
955        let value =
956            truncated_normal::<_, f32>(array!([0.0, 0.5]), array!([0.5, 1.0]), None, &key).unwrap();
957        let expected = Array::from_slice(&[0.10, 0.88], &[2]);
958
959        assert_array_eq(
960            value,
961            expected,
962            tolerances::ROUNDED_TWO_DECIMALS.rtol,
963            tolerances::ROUNDED_TWO_DECIMALS.atol,
964        );
965    }
966
967    #[test]
968    fn test_gumbel() {
969        let key = key(0).unwrap();
970        let value = gumbel::<f32>(None, &key).unwrap();
971        assert_array_eq(
972            value,
973            Array::from_f32(0.13),
974            tolerances::ROUNDED_TWO_DECIMALS.rtol,
975            tolerances::ROUNDED_TWO_DECIMALS.atol,
976        );
977    }
978
979    #[test]
980    fn test_logits() {
981        let key = key(0).unwrap();
982        let logits = Array::zeros::<u32>(&[5, 20]).unwrap();
983        let result = categorical(&logits, None, None, &key).unwrap();
984
985        assert_eq!(result.shape(), [5]);
986
987        let expected = Array::from_slice(&[1_u32, 1, 17, 17, 17], &[5]);
988        assert_array_eq_with_context(
989            result,
990            expected,
991            tolerances::EXACT.rtol,
992            tolerances::EXACT.atol,
993            "categorical default sample values",
994        );
995    }
996
997    #[test]
998    fn test_logits_count() {
999        let key = key(0).unwrap();
1000        let logits = Array::zeros::<u32>(&[5, 20]).unwrap();
1001        let result = categorical(&logits, None, ShapeOrCount::Count(2), &key).unwrap();
1002
1003        assert_eq!(result.shape(), [5, 2]);
1004
1005        let expected = Array::from_slice(&[16_u32, 3, 14, 10, 17, 7, 6, 8, 12, 8], &[5, 2]);
1006        assert_array_eq_with_context(
1007            result,
1008            expected,
1009            tolerances::EXACT.rtol,
1010            tolerances::EXACT.atol,
1011            "categorical counted sample values",
1012        );
1013    }
1014
1015    #[test]
1016    fn test_random_state_new() {
1017        let state = RandomState::new().unwrap();
1018        assert_eq!(state.as_array().shape(), &[2]);
1019    }
1020
1021    #[test]
1022    fn test_random_state_with_seed_deterministic() {
1023        let s1 = RandomState::with_seed(42).unwrap();
1024        let s2 = RandomState::with_seed(42).unwrap();
1025        assert_array_eq(
1026            s1.as_array(),
1027            s2.as_array(),
1028            tolerances::EXACT.rtol,
1029            tolerances::EXACT.atol,
1030        );
1031    }
1032
1033    #[test]
1034    fn test_random_state_next_key_advances() {
1035        let mut state = RandomState::with_seed(0).unwrap();
1036        let k1 = state.next_key().unwrap();
1037        let k2 = state.next_key().unwrap();
1038        assert!(!k1.eq_exact(&k2).unwrap());
1039    }
1040
1041    #[test]
1042    fn test_random_state_from_key_roundtrip() {
1043        let original = RandomState::with_seed(99).unwrap();
1044        let arr = original.as_array().clone();
1045        let restored = RandomState::from_key(arr);
1046        assert_array_eq(
1047            original.as_array(),
1048            restored.as_array(),
1049            tolerances::EXACT.rtol,
1050            tolerances::EXACT.atol,
1051        );
1052    }
1053
1054    #[test]
1055    fn test_random_state_updatable() {
1056        use crate::utils::Updatable;
1057        let mut state = RandomState::with_seed(0).unwrap();
1058        let projection = state.state_projection().unwrap();
1059        assert_eq!(projection.len(), 1);
1060        assert_eq!(projection.values().count(), 1);
1061    }
1062
1063    #[test]
1064    fn test_random_state_default() {
1065        let state = RandomState::default();
1066        assert_eq!(state.as_array().shape(), &[2]);
1067    }
1068
1069    #[test]
1070    fn test_random_seed_same() {
1071        // Same random seed should produce the same results
1072        let seed = 23;
1073        let mut results = Vec::new();
1074        let f = || {
1075            let sum = uniform::<_, f32>(0.0, 1.0, &[10, 10], None)?.sum(None)?;
1076            Ok::<_, crate::error::Exception>(sum.item_exact::<f32>())
1077        };
1078        for _ in 0..10 {
1079            let mut state = RandomState::new().unwrap();
1080            state.seed(seed).unwrap();
1081            let result = with_random_state(state, f).unwrap();
1082            results.push(result);
1083        }
1084
1085        // Check that all results are the same within a small tolerance
1086        let first = results[0];
1087        for result in &results[1..] {
1088            assert_float_eq!(
1089                first,
1090                *result,
1091                abs <= 0.01,
1092                "Results should be equal for the same seed"
1093            );
1094        }
1095    }
1096}
mlx_rs/random.rs

mlx_rs/
random.rs
