safemlx/
lib.rs

1//! Safe, low-level Rust bindings to MLX.
2//!
3//! `safemlx` exposes MLX arrays, devices, streams, native operators, function
4//! transforms, and accelerator/runtime facilities. It intentionally does not
5//! define neural-network layers, models, optimizers, checkpoint policy, or
6//! other framework abstractions; those belong to consumers such as
7//! `eredu-backend-mlx`.
8//!
9//! Operations are lazy and are scheduled on an explicit [`Stream`]. Call
10//! [`Array::evaluated`] (or a function in [`transforms`]) to submit work before
11//! accessing host data. Native MLX `.npy` and `safetensors` I/O is exposed on
12//! [`Array`]; higher-level checkpoint formats are outside this crate.
13
14#![deny(unused_unsafe, missing_debug_implementations, missing_docs)]
15#![cfg_attr(test, allow(clippy::approx_constant))]
16
17#[macro_use]
18pub mod macros;
19
20mod array;
21#[cfg(feature = "cuda")]
22pub mod cuda;
23mod device;
24pub mod distributed;
25mod dtype;
26pub mod error;
27mod event;
28pub mod fast;
29pub mod fft;
30mod host_transfer;
31pub mod linalg;
32pub mod memory;
33#[cfg(feature = "metal")]
34pub mod metal;
35pub mod ops;
36pub mod random;
37mod stream;
38pub mod system;
39pub mod transforms;
40pub mod utils;
41
42pub use array::*;
43pub use device::*;
44pub use dtype::*;
45pub use event::*;
46pub use host_transfer::*;
47pub use stream::*;
48
49#[cfg(test)]
50pub(crate) fn test_stream() -> &'static Stream {
51    Box::leak(Box::new(Stream::new_with_device(&Device::new(
52        DeviceType::Cpu,
53        0,
54    ))))
55}
56
57#[cfg(test)]
58pub(crate) fn test_key(seed: u64, stream: &Stream) -> Array {
59    use crate::ops::indexing::TryIndexOp;
60
61    random::split_n(random::key(seed).unwrap(), 2, stream)
62        .unwrap()
63        .try_index_device(1, stream)
64        .unwrap()
65}
66
67#[cfg(test)]
68pub(crate) fn test_concurrency() -> usize {
69    std::thread::available_parallelism()
70        .map(|parallelism| parallelism.get())
71        .unwrap_or(2)
72        .clamp(2, 16)
73}
74
75pub(crate) mod constants {
76    pub(crate) const DEFAULT_STACK_VEC_LEN: usize = 4;
77}
78
79pub(crate) mod sealed {
80    pub trait Sealed {}
81
82    impl Sealed for () {}
83    impl<A> Sealed for (A,) where A: Sealed {}
84    impl<A, B> Sealed for (A, B)
85    where
86        A: Sealed,
87        B: Sealed,
88    {
89    }
90    impl<A, B, C> Sealed for (A, B, C)
91    where
92        A: Sealed,
93        B: Sealed,
94        C: Sealed,
95    {
96    }
97}
safemlx/lib.rs

safemlx/
lib.rs
