mlx_rs/ops/
io.rs

1use crate::error::IoError;
2use crate::utils::guard::Guarded;
3use crate::utils::io::SafeTensors;
4use crate::utils::SUCCESS;
5use crate::{Array, Stream};
6use std::collections::HashMap;
7use std::ffi::CString;
8use std::path::Path;
9fn check_file_extension(path: &Path, expected: &str) -> Result<(), IoError> {
10    match path.extension().and_then(|ext| ext.to_str()) {
11        Some(ext) if ext == expected => Ok(()),
12        _ => Err(IoError::UnsupportedFormat),
13    }
14}
15
16impl Array {
17    /// Load array from a binary file in `.npy` format.
18    ///
19    /// # Params
20    ///
21    /// - path: path of file to load
22    /// - stream: stream or device to evaluate on
23    pub fn load_numpy(path: impl AsRef<Path>) -> Result<Array, IoError> {
24        let stream = Stream::thread_local_or_cpu();
25        let path = path.as_ref();
26        if !path.is_file() {
27            return Err(IoError::NotFile);
28        }
29        let c_path = CString::new(path.to_str().ok_or(IoError::InvalidUtf8)?)?;
30        check_file_extension(path, "npy")?;
31
32        Array::try_from_op(|res| unsafe {
33            mlx_sys::mlx_load(res, c_path.as_ptr(), stream.as_ref().as_ptr())
34        })
35        .map_err(Into::into)
36    }
37
38    /// Compatibility shim for [`load_numpy`].
39    #[deprecated(
40        since = "0.26.0",
41        note = "use `with_stream` or `with_device` around `load_numpy`"
42    )]
43    pub fn load_numpy_device(
44        path: impl AsRef<Path>,
45        stream: impl AsRef<Stream>,
46    ) -> Result<Array, IoError> {
47        crate::with_stream(stream.as_ref(), || Self::load_numpy(path))
48    }
49
50    /// Load dictionary of ``MLXArray`` from a `safetensors` file.
51    ///
52    /// # Params
53    ///
54    /// - path: path of file to load
55    /// - stream: stream or device to evaluate on
56    ///
57    pub fn load_safetensors(path: impl AsRef<Path>) -> Result<HashMap<String, Array>, IoError> {
58        let stream = Stream::thread_local_or_cpu();
59        let safetensors = SafeTensors::load_device(path.as_ref(), stream)?;
60        let data = safetensors.data()?;
61        Ok(data)
62    }
63
64    /// Compatibility shim for [`load_safetensors`].
65    #[deprecated(
66        since = "0.26.0",
67        note = "use `with_stream` or `with_device` around `load_safetensors`"
68    )]
69    pub fn load_safetensors_device(
70        path: impl AsRef<Path>,
71        stream: impl AsRef<Stream>,
72    ) -> Result<HashMap<String, Array>, IoError> {
73        crate::with_stream(stream.as_ref(), || Self::load_safetensors(path))
74    }
75
76    /// Load dictionary of ``MLXArray`` and metadata `[String:String]` from a `safetensors` file.
77    ///
78    /// # Params
79    ///
80    /// - path: path of file to load
81    /// - stream: stream or device to evaluate on
82    #[allow(clippy::type_complexity)]
83    pub fn load_safetensors_with_metadata(
84        path: impl AsRef<Path>,
85    ) -> Result<(HashMap<String, Array>, HashMap<String, String>), IoError> {
86        let stream = Stream::thread_local_or_cpu();
87        let safetensors = SafeTensors::load_device(path.as_ref(), stream)?;
88        let data = safetensors.data()?;
89        let metadata = safetensors.metadata()?;
90
91        Ok((data, metadata))
92    }
93
94    /// Compatibility shim for [`load_safetensors_with_metadata`].
95    #[allow(clippy::type_complexity)]
96    #[deprecated(
97        since = "0.26.0",
98        note = "use `with_stream` or `with_device` around `load_safetensors_with_metadata`"
99    )]
100    pub fn load_safetensors_with_metadata_device(
101        path: impl AsRef<Path>,
102        stream: impl AsRef<Stream>,
103    ) -> Result<(HashMap<String, Array>, HashMap<String, String>), IoError> {
104        crate::with_stream(stream.as_ref(), || {
105            Self::load_safetensors_with_metadata(path)
106        })
107    }
108
109    /// Save array to a binary file in `.npy`format.
110    ///
111    /// # Params
112    ///
113    /// - array: array to save
114    /// - url: URL of file to load
115    pub fn save_numpy(&self, path: impl AsRef<Path>) -> Result<(), IoError> {
116        let path = path.as_ref();
117        check_file_extension(path, "npy")?;
118        let c_path = CString::new(path.to_str().ok_or(IoError::InvalidUtf8)?)?;
119
120        unsafe { mlx_sys::mlx_save(c_path.as_ptr(), self.as_ptr()) };
121
122        Ok(())
123    }
124
125    /// Save dictionary of arrays in `safetensors` format.
126    ///
127    /// # Params
128    ///
129    /// - arrays: arrays to save
130    /// - metadata: metadata to save
131    /// - path: path of file to save
132    /// - stream: stream or device to evaluate on
133    pub fn save_safetensors<'a, I, S, V>(
134        arrays: I,
135        metadata: impl Into<Option<&'a HashMap<String, String>>>,
136        path: impl AsRef<Path>,
137    ) -> Result<(), IoError>
138    where
139        I: IntoIterator<Item = (S, V)>,
140        S: AsRef<str>,
141        V: AsRef<Array>,
142    {
143        crate::error::INIT_ERR_HANDLER
144            .with(|init| init.call_once(crate::error::setup_mlx_error_handler));
145
146        let path = path.as_ref();
147
148        check_file_extension(path, "safetensors")?;
149
150        let arrays = unsafe {
151            let data = mlx_sys::mlx_map_string_to_array_new();
152            for (key, array) in arrays.into_iter() {
153                let key = CString::new(key.as_ref())?;
154
155                let status = mlx_sys::mlx_map_string_to_array_insert(
156                    data,
157                    key.as_ptr(),
158                    array.as_ref().as_ptr(),
159                );
160
161                if status != SUCCESS {
162                    mlx_sys::mlx_map_string_to_array_free(data);
163                    return Err(crate::error::exception_from_status(
164                        status,
165                        "inserting a safetensors array",
166                    )
167                    .into());
168                }
169            }
170            data
171        };
172
173        let default_metadata = HashMap::new();
174        let metadata_ref = metadata.into().unwrap_or(&default_metadata);
175
176        let metadata = unsafe {
177            let data = mlx_sys::mlx_map_string_to_string_new();
178            for (key, value) in metadata_ref.iter() {
179                let key = CString::new(key.as_str())?;
180                let value = CString::new(value.as_str())?;
181
182                let status =
183                    mlx_sys::mlx_map_string_to_string_insert(data, key.as_ptr(), value.as_ptr());
184
185                if status != SUCCESS {
186                    mlx_sys::mlx_map_string_to_string_free(data);
187                    return Err(crate::error::exception_from_status(
188                        status,
189                        "inserting safetensors metadata",
190                    )
191                    .into());
192                }
193            }
194            data
195        };
196
197        let c_path = CString::new(path.to_str().ok_or(IoError::InvalidUtf8)?)?;
198
199        unsafe {
200            let status = mlx_sys::mlx_save_safetensors(c_path.as_ptr(), arrays, metadata);
201
202            let last_error = match status {
203                SUCCESS => None,
204                _ => Some(crate::error::exception_from_status(
205                    status,
206                    "saving safetensors",
207                )),
208            };
209
210            mlx_sys::mlx_map_string_to_array_free(arrays);
211            mlx_sys::mlx_map_string_to_string_free(metadata);
212
213            if let Some(error) = last_error {
214                return Err(error.into());
215            }
216        };
217
218        Ok(())
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use crate::Array;
225
226    #[test]
227    fn test_save_arrays() {
228        let tmp_dir = tempfile::tempdir().unwrap();
229        let path = tmp_dir.path().join("test.safetensors");
230
231        let mut arrays = std::collections::HashMap::new();
232        arrays.insert("foo".to_string(), Array::ones::<i32>(&[1, 2]).unwrap());
233        arrays.insert("bar".to_string(), Array::zeros::<i32>(&[2, 1]).unwrap());
234
235        Array::save_safetensors(&arrays, None, &path).unwrap();
236
237        let loaded_arrays = Array::load_safetensors(&path).unwrap();
238
239        // compare values
240        let mut loaded_keys: Vec<_> = loaded_arrays.keys().cloned().collect();
241        let mut original_keys: Vec<_> = arrays.keys().cloned().collect();
242        loaded_keys.sort();
243        original_keys.sort();
244        assert_eq!(loaded_keys, original_keys);
245
246        for key in loaded_keys {
247            let loaded_array = loaded_arrays.get(&key).unwrap();
248            let original_array = arrays.get(&key).unwrap();
249            assert!(loaded_array
250                .all_close(original_array, None, None, None)
251                .unwrap());
252        }
253    }
254
255    #[test]
256    fn test_save_array() {
257        let tmp_dir = tempfile::tempdir().unwrap();
258        let path = tmp_dir.path().join("test.npy");
259
260        let a = Array::ones::<i32>(&[2, 4]).unwrap();
261        a.save_numpy(&path).unwrap();
262
263        let b = Array::load_numpy(&path).unwrap();
264        assert!(a.all_close(&b, None, None, None).unwrap());
265    }
266}
mlx_rs/ops/io.rs

mlx_rs/ops/
io.rs
