Struct RandomState 

pub struct RandomState { /* private fields */ }

Random state for reproducible random number generation.

This struct holds the PRNG state and can be used with compiled functions to properly track random state across JIT compilation boundaries.

Compilation Support

RandomState implements Updatable, making it compatible with compile_with_state. This is the Rust equivalent of Python’s @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state).

Example

use mlx_rs::random::RandomState;
use mlx_rs::transforms::compile::compile_with_state;
use mlx_rs::random::categorical;
use mlx_rs::Array;

let mut state = RandomState::with_seed(42).unwrap();
let logits = Array::zeros::<f32>(&[1, 10]).unwrap();
let mut compiled = compile_with_state(
    |state: &mut RandomState, x: &Array| {
        let key = state.next_key()?;
        categorical(x, None, None, Some(&key))
    },
    None
);
let result = compiled(&mut state, &logits).unwrap();

Implementations

impl RandomState

pub fn new() -> Result<Self>

Create a new random state with a time-based seed.

pub fn with_seed(seed: u64) -> Result<Self>

Create a new random state from a specific seed.

Use this for reproducible random number generation.

pub fn from_key(key: Array) -> Self

Create a random state from an existing key array.

The key must be a valid PRNG key (typically created via random::key()).

pub fn next_key(&mut self) -> Result<Array>

Get the next random key, advancing the state.

This splits the current state into two keys: one becomes the new state, and the other is returned for use in random operations.

pub fn seed(&mut self, seed: u64) -> Result<()>

Reseed the random state.

pub fn as_array(&self) -> &Array

Get a reference to the underlying state array.

This is useful for inspection or manual state management.

pub fn as_array_mut(&mut self) -> &mut Array

Get a mutable reference to the underlying state array.

Note

Modifying the state array directly may break the PRNG invariants. Prefer using seed() or next_key() instead.

Trait Implementations

impl Clone for RandomState

fn clone(&self) -> RandomState

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for RandomState

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for RandomState

fn default() -> Self

Creates a new RandomState with a time-based seed.

Panics

Panics if the underlying PRNG key creation fails, which should not occur under normal conditions.

impl Updatable for RandomState

fn state_projection( &mut self, ) -> Result<StateProjection<'_>, StateProjectionError>

Declare all required and optional state slots.