mlx-rs

Unlike Python's MLX, mlx-rs requires you to be explicit about which arrays are traced in the computation graph. To avoid segfaults, do not capture arrays from the outer scope inside your loss function closure. Instead, pass all required arrays as inputs to the closure.

Correct Pattern

1. Define a loss function that accepts a slice of arrays: |inputs: &[Array]| -> Result<Array, Exception>.
2. Extract the necessary arrays (weights, data, targets) from the inputs slice inside the closure.
3. Use transforms::grad to create the gradient function.
4. Pass all required arrays in the input slice to the gradient function.
5. When training, manually update the arrays in your input vector and call .eval() on the updated array.

// 1. Define loss function with explicit inputs
let loss_fn = |inputs: &[Array]| -> Result<Array, Exception> {
    let w = &inputs[0];
    let x = &inputs[1];
    let y = &inputs[2];

    let y_pred = x.matmul(w)?;
    let loss = Array::from_f32(0.5) * ops::mean(&ops::square(y_pred - y)?, None, None)?;
    Ok(loss)
};

let argnums = &[0]; // Differentiate with respect to the first argument (w)
let mut inputs = vec![w, x, y];

// 2. Compute gradients
let grad = transforms::grad(loss_fn, argnums)(&inputs)?;

// 3. Update weights in the training loop
inputs[0] = &inputs[0] - Array::from_f32(learning_rate) * grad;
inputs[0].eval()?; // Ensure the update is materialized

In MLX, performing operations on arrays does not trigger immediate computation; instead, it records a compute graph. To execute the operations and compute the actual values, you must explicitly call Array::eval or trigger an implicit evaluation.

Explicit Evaluation

Use Array::eval() to evaluate the output of an operation.

Implicit Evaluation

An array is automatically evaluated when you:

• Use Array::item to get a scalar value.
• Print the array with println!.
• Access the underlying data via Array::as_slice.
• Save the array using Array::save_numpy or Array::save_safetensors.

Best Practices

• Batching: Avoid calling eval() too frequently (e.g., inside every single arithmetic operation) as there is fixed overhead per evaluation. A good pattern is to evaluate at the end of each training iteration (e.g., after computing loss and gradients).
• Control Flow: Using scalar arrays for control flow (e.g., if array.item() > 0.5 { ... }) triggers an evaluation, which can be inefficient if done frequently.

use mlx_rs::{array, transforms::eval};

let a = array!([1, 2, 3, 4]);
let b = array!([1.0, 2.0, 3.0, 4.0]);

let c = &a + &b; // c is not evaluated
c.eval().unwrap(); // evaluates c

let d = &a + &b;
println!("{:?}", d); // evaluates d

let e = &a + &b;
let e_slice: &[f32] = e.as_slice(); // evaluates e

RandomState is used for reproducible random number generation. It holds the PRNG state and is compatible with compile_with_state, allowing random state to be tracked across JIT compilation boundaries (similar to Python's @partial(mx.compile, inputs=mx.random.state, ...)).

Key Methods

• RandomState::new(): Creates a new state using a time-based seed.
• RandomState::with_seed(seed: u64): Creates a state from a specific seed for reproducibility.
• RandomState::from_key(key: Array): Creates a state from an existing PRNG key array.
• RandomState::next_key(): Advances the state and returns a new key for use in random operations.
• RandomState::seed(seed: u64): Reseeds the existing state.
• RandomState::as_array(): Returns a reference to the underlying state array.

Usage with Compiled Functions

To use RandomState within a compiled function, pass it as a mutable argument to compile_with_state so the state can be updated during execution.

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();