Function with_stream 

pub fn with_stream<F, T>(stream: &Stream, f: F) -> T
where
    F: FnOnce() -> T,

Uses stream for operations constructed during f.

Scopes are synchronous, thread-local, nestable, and restore the previous stream if f panics. To select a stream for one operation, put only that operation in the closure:

use mlx_rs::{with_stream, Array, Stream};

let input = Array::from_slice(&[1.0_f32, 2.0, 3.0, 4.0], &[4]);
let stream = Stream::cpu();
let output = with_stream(&stream, || mlx_rs::fft::fft(&input, None, None)).unwrap();
assert_eq!(output.shape(), &[4]);