mlx_rs/ops/
conversion.rs

1use mlx_internal_macros::generate_macro;
2
3use crate::{error::Result, utils::guard::Guarded, Array, ArrayElement, Dtype, Stream};
4
5impl Array {
6    /// Convert an array to FP8 (E4M3) format.
7    ///
8    /// The input array must be a floating point type (float32, float16, or bfloat16).
9    /// Values outside the representable range of FP8 E4M3 (-448 to 448) will be clipped.
10    ///
11    /// # Returns
12    ///
13    /// An array with dtype uint8 containing the FP8 E4M3 encoded values.
14    pub fn to_fp8(&self) -> Result<Array> {
15        let stream = Stream::thread_local_or_default();
16        Array::try_from_op(|res| unsafe {
17            mlx_sys::mlx_to_fp8(res, self.as_ptr(), stream.as_ref().as_ptr())
18        })
19    }
20
21    /// Compatibility shim for [`to_fp8`].
22    #[deprecated(
23        since = "0.26.0",
24        note = "use `with_stream` or `with_device` around `to_fp8`"
25    )]
26    pub fn to_fp8_device(&self, stream: impl AsRef<Stream>) -> Result<Array> {
27        crate::with_stream(stream.as_ref(), || self.to_fp8())
28    }
29
30    /// Convert an FP8 (E4M3) encoded array back to a floating point type.
31    ///
32    /// The input array should be a uint8 array containing FP8 E4M3 encoded values.
33    ///
34    /// # Params
35    ///
36    /// - `dtype`: The target floating point dtype (float32, float16, or bfloat16)
37    pub fn from_fp8(&self, dtype: Dtype) -> Result<Array> {
38        let stream = Stream::thread_local_or_default();
39        Array::try_from_op(|res| unsafe {
40            mlx_sys::mlx_from_fp8(res, self.as_ptr(), dtype.into(), stream.as_ref().as_ptr())
41        })
42    }
43
44    /// Compatibility shim for [`from_fp8`].
45    #[deprecated(
46        since = "0.26.0",
47        note = "use `with_stream` or `with_device` around `from_fp8`"
48    )]
49    pub fn from_fp8_device(&self, dtype: Dtype, stream: impl AsRef<Stream>) -> Result<Array> {
50        crate::with_stream(stream.as_ref(), || self.from_fp8(dtype))
51    }
52
53    /// Create a new array with the contents converted to the given [ArrayElement] type.
54    ///
55    /// # Example
56    ///
57    /// ```rust
58    /// use mlx_rs::{Array, Dtype};
59    ///
60    /// let array = Array::from_slice(&[1i16,2,3], &[3]);
61    /// let mut new_array = array.as_type::<f32>().unwrap();
62    ///
63    /// assert_eq!(new_array.dtype(), Dtype::Float32);
64    /// assert_eq!(new_array.shape(), &[3]);
65    /// assert_eq!(new_array.item_size(), 4);
66    /// assert_eq!(new_array.as_slice::<f32>(), &[1.0,2.0,3.0]);
67    /// ```
68    pub fn as_type<T: ArrayElement>(&self) -> Result<Array> {
69        self.as_dtype(T::DTYPE)
70    }
71
72    /// Compatibility shim for [`as_type`].
73    #[deprecated(
74        since = "0.26.0",
75        note = "use `with_stream` or `with_device` around `as_type`"
76    )]
77    pub fn as_type_device<T: ArrayElement>(&self, stream: impl AsRef<Stream>) -> Result<Array> {
78        crate::with_stream(stream.as_ref(), || self.as_type::<T>())
79    }
80
81    /// Same as `as_type` but with a [`Dtype`] argument.
82    pub fn as_dtype(&self, dtype: Dtype) -> Result<Array> {
83        let stream = Stream::thread_local_or_default();
84        Array::try_from_op(|res| unsafe {
85            mlx_sys::mlx_astype(res, self.as_ptr(), dtype.into(), stream.as_ref().as_ptr())
86        })
87    }
88
89    /// Compatibility shim for [`as_dtype`].
90    #[deprecated(
91        since = "0.26.0",
92        note = "use `with_stream` or `with_device` around `as_dtype`"
93    )]
94    pub fn as_dtype_device(&self, dtype: Dtype, stream: impl AsRef<Stream>) -> Result<Array> {
95        crate::with_stream(stream.as_ref(), || self.as_dtype(dtype))
96    }
97
98    /// View the array as a different type.
99    ///
100    /// The output array will change along the last axis if the input array's
101    /// type and the output array's type do not have the same size.
102    ///
103    /// _Note: the view op does not imply that the input and output arrays share
104    /// their underlying data. The view only guarantees that the binary
105    /// representation of each element (or group of elements) is the same._
106    ///
107    pub fn view<T: ArrayElement>(&self) -> Result<Array> {
108        self.view_dtype(T::DTYPE)
109    }
110
111    /// Compatibility shim for [`view`].
112    #[deprecated(
113        since = "0.26.0",
114        note = "use `with_stream` or `with_device` around `view`"
115    )]
116    pub fn view_device<T: ArrayElement>(&self, stream: impl AsRef<Stream>) -> Result<Array> {
117        crate::with_stream(stream.as_ref(), || self.view::<T>())
118    }
119
120    /// Same as `view` but with a [`Dtype`] argument.
121    pub fn view_dtype(&self, dtype: Dtype) -> Result<Array> {
122        let stream = Stream::thread_local_or_default();
123        Array::try_from_op(|res| unsafe {
124            mlx_sys::mlx_view(res, self.as_ptr(), dtype.into(), stream.as_ref().as_ptr())
125        })
126    }
127
128    /// Compatibility shim for [`view_dtype`].
129    #[deprecated(
130        since = "0.26.0",
131        note = "use `with_stream` or `with_device` around `view_dtype`"
132    )]
133    pub fn view_dtype_device(&self, dtype: Dtype, stream: impl AsRef<Stream>) -> Result<Array> {
134        crate::with_stream(stream.as_ref(), || self.view_dtype(dtype))
135    }
136}
137
138/// Convert an array to FP8 (E4M3) format.
139///
140/// See [`Array::to_fp8`] for more details.
141pub fn to_fp8(a: impl AsRef<Array>) -> Result<Array> {
142    a.as_ref().to_fp8()
143}
144
145/// Compatibility shim for [`to_fp8`].
146#[generate_macro(customize(forwarding_shim = true))]
147#[deprecated(
148    since = "0.26.0",
149    note = "use `with_stream` or `with_device` around `to_fp8`"
150)]
151pub fn to_fp8_device(
152    a: impl AsRef<Array>,
153    #[optional] stream: impl AsRef<Stream>,
154) -> Result<Array> {
155    crate::with_stream(stream.as_ref(), || to_fp8(a))
156}
157
158/// Convert an FP8 (E4M3) encoded array back to a floating point type.
159///
160/// See [`Array::from_fp8`] for more details.
161pub fn from_fp8(a: impl AsRef<Array>, dtype: Dtype) -> Result<Array> {
162    a.as_ref().from_fp8(dtype)
163}
164
165/// Compatibility shim for [`from_fp8`].
166#[generate_macro(customize(forwarding_shim = true))]
167#[deprecated(
168    since = "0.26.0",
169    note = "use `with_stream` or `with_device` around `from_fp8`"
170)]
171pub fn from_fp8_device(
172    a: impl AsRef<Array>,
173    dtype: Dtype,
174    #[optional] stream: impl AsRef<Stream>,
175) -> Result<Array> {
176    crate::with_stream(stream.as_ref(), || from_fp8(a, dtype))
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182    use crate::complex64;
183    use half::{bf16, f16};
184    use pretty_assertions::assert_eq;
185
186    macro_rules! test_as_type {
187        ($src_type:ty, $src_val:expr, $dst_type:ty, $dst_val:expr, $len:expr) => {
188            paste::paste! {
189                #[test]
190                fn [<test_as_type_ $src_type _ $dst_type>]() {
191                    let array = Array::from_slice(&[$src_val; $len], &[$len as i32]);
192                    let new_array = array.as_type::<$dst_type>().unwrap();
193
194                    assert_eq!(new_array.dtype(), $dst_type::DTYPE);
195                    assert_eq!(new_array.shape(), &[3]);
196                    assert_eq!(new_array.item_size(), std::mem::size_of::<$dst_type>());
197                    assert_eq!(new_array.as_slice::<$dst_type>(), &[$dst_val; $len]);
198                }
199            }
200        };
201    }
202
203    test_as_type!(bool, true, i8, 1, 3);
204    test_as_type!(bool, true, i16, 1, 3);
205    test_as_type!(bool, true, i32, 1, 3);
206    test_as_type!(bool, true, i64, 1, 3);
207    test_as_type!(bool, true, u8, 1, 3);
208    test_as_type!(bool, true, u16, 1, 3);
209    test_as_type!(bool, true, u32, 1, 3);
210    test_as_type!(bool, true, u64, 1, 3);
211    test_as_type!(bool, true, f32, 1.0, 3);
212    test_as_type!(bool, true, f16, f16::from_f32(1.0), 3);
213    test_as_type!(bool, true, bf16, bf16::from_f32(1.0), 3);
214    test_as_type!(bool, true, complex64, complex64::new(1.0, 0.0), 3);
215
216    test_as_type!(i8, 1, bool, true, 3);
217    test_as_type!(i8, 1, i16, 1, 3);
218    test_as_type!(i8, 1, i32, 1, 3);
219    test_as_type!(i8, 1, i64, 1, 3);
220    test_as_type!(i8, 1, u8, 1, 3);
221    test_as_type!(i8, 1, u16, 1, 3);
222    test_as_type!(i8, 1, u32, 1, 3);
223    test_as_type!(i8, 1, u64, 1, 3);
224    test_as_type!(i8, 1, f32, 1.0, 3);
225    test_as_type!(i8, 1, f16, f16::from_f32(1.0), 3);
226    test_as_type!(i8, 1, bf16, bf16::from_f32(1.0), 3);
227    test_as_type!(i8, 1, complex64, complex64::new(1.0, 0.0), 3);
228
229    test_as_type!(i16, 1, bool, true, 3);
230    test_as_type!(i16, 1, i8, 1, 3);
231    test_as_type!(i16, 1, i32, 1, 3);
232    test_as_type!(i16, 1, i64, 1, 3);
233    test_as_type!(i16, 1, u8, 1, 3);
234    test_as_type!(i16, 1, u16, 1, 3);
235    test_as_type!(i16, 1, u32, 1, 3);
236    test_as_type!(i16, 1, u64, 1, 3);
237    test_as_type!(i16, 1, f32, 1.0, 3);
238    test_as_type!(i16, 1, f16, f16::from_f32(1.0), 3);
239    test_as_type!(i16, 1, bf16, bf16::from_f32(1.0), 3);
240    test_as_type!(i16, 1, complex64, complex64::new(1.0, 0.0), 3);
241
242    test_as_type!(i32, 1, bool, true, 3);
243    test_as_type!(i32, 1, i8, 1, 3);
244    test_as_type!(i32, 1, i16, 1, 3);
245    test_as_type!(i32, 1, i64, 1, 3);
246    test_as_type!(i32, 1, u8, 1, 3);
247    test_as_type!(i32, 1, u16, 1, 3);
248    test_as_type!(i32, 1, u32, 1, 3);
249    test_as_type!(i32, 1, u64, 1, 3);
250    test_as_type!(i32, 1, f32, 1.0, 3);
251    test_as_type!(i32, 1, f16, f16::from_f32(1.0), 3);
252    test_as_type!(i32, 1, bf16, bf16::from_f32(1.0), 3);
253    test_as_type!(i32, 1, complex64, complex64::new(1.0, 0.0), 3);
254
255    test_as_type!(i64, 1, bool, true, 3);
256    test_as_type!(i64, 1, i8, 1, 3);
257    test_as_type!(i64, 1, i16, 1, 3);
258    test_as_type!(i64, 1, i32, 1, 3);
259    test_as_type!(i64, 1, u8, 1, 3);
260    test_as_type!(i64, 1, u16, 1, 3);
261    test_as_type!(i64, 1, u32, 1, 3);
262    test_as_type!(i64, 1, u64, 1, 3);
263    test_as_type!(i64, 1, f32, 1.0, 3);
264    test_as_type!(i64, 1, f16, f16::from_f32(1.0), 3);
265    test_as_type!(i64, 1, bf16, bf16::from_f32(1.0), 3);
266    test_as_type!(i64, 1, complex64, complex64::new(1.0, 0.0), 3);
267
268    test_as_type!(u8, 1, bool, true, 3);
269    test_as_type!(u8, 1, i8, 1, 3);
270    test_as_type!(u8, 1, i16, 1, 3);
271    test_as_type!(u8, 1, i32, 1, 3);
272    test_as_type!(u8, 1, i64, 1, 3);
273    test_as_type!(u8, 1, u16, 1, 3);
274    test_as_type!(u8, 1, u32, 1, 3);
275    test_as_type!(u8, 1, u64, 1, 3);
276    test_as_type!(u8, 1, f32, 1.0, 3);
277    test_as_type!(u8, 1, f16, f16::from_f32(1.0), 3);
278    test_as_type!(u8, 1, bf16, bf16::from_f32(1.0), 3);
279    test_as_type!(u8, 1, complex64, complex64::new(1.0, 0.0), 3);
280
281    test_as_type!(u16, 1, bool, true, 3);
282    test_as_type!(u16, 1, i8, 1, 3);
283    test_as_type!(u16, 1, i16, 1, 3);
284    test_as_type!(u16, 1, i32, 1, 3);
285    test_as_type!(u16, 1, i64, 1, 3);
286    test_as_type!(u16, 1, u8, 1, 3);
287    test_as_type!(u16, 1, u32, 1, 3);
288    test_as_type!(u16, 1, u64, 1, 3);
289    test_as_type!(u16, 1, f32, 1.0, 3);
290    test_as_type!(u16, 1, f16, f16::from_f32(1.0), 3);
291    test_as_type!(u16, 1, bf16, bf16::from_f32(1.0), 3);
292    test_as_type!(u16, 1, complex64, complex64::new(1.0, 0.0), 3);
293
294    test_as_type!(u32, 1, bool, true, 3);
295    test_as_type!(u32, 1, i8, 1, 3);
296    test_as_type!(u32, 1, i16, 1, 3);
297    test_as_type!(u32, 1, i32, 1, 3);
298    test_as_type!(u32, 1, i64, 1, 3);
299    test_as_type!(u32, 1, u8, 1, 3);
300    test_as_type!(u32, 1, u16, 1, 3);
301    test_as_type!(u32, 1, u64, 1, 3);
302    test_as_type!(u32, 1, f32, 1.0, 3);
303    test_as_type!(u32, 1, f16, f16::from_f32(1.0), 3);
304    test_as_type!(u32, 1, bf16, bf16::from_f32(1.0), 3);
305    test_as_type!(u32, 1, complex64, complex64::new(1.0, 0.0), 3);
306
307    test_as_type!(u64, 1, bool, true, 3);
308    test_as_type!(u64, 1, i8, 1, 3);
309    test_as_type!(u64, 1, i16, 1, 3);
310    test_as_type!(u64, 1, i32, 1, 3);
311    test_as_type!(u64, 1, i64, 1, 3);
312    test_as_type!(u64, 1, u8, 1, 3);
313    test_as_type!(u64, 1, u16, 1, 3);
314    test_as_type!(u64, 1, u32, 1, 3);
315    test_as_type!(u64, 1, f32, 1.0, 3);
316    test_as_type!(u64, 1, f16, f16::from_f32(1.0), 3);
317    test_as_type!(u64, 1, bf16, bf16::from_f32(1.0), 3);
318    test_as_type!(u64, 1, complex64, complex64::new(1.0, 0.0), 3);
319
320    test_as_type!(f32, 1.0, bool, true, 3);
321    test_as_type!(f32, 1.0, i8, 1, 3);
322    test_as_type!(f32, 1.0, i16, 1, 3);
323    test_as_type!(f32, 1.0, i32, 1, 3);
324    test_as_type!(f32, 1.0, i64, 1, 3);
325    test_as_type!(f32, 1.0, u8, 1, 3);
326    test_as_type!(f32, 1.0, u16, 1, 3);
327    test_as_type!(f32, 1.0, u32, 1, 3);
328    test_as_type!(f32, 1.0, u64, 1, 3);
329    test_as_type!(f32, 1.0, f16, f16::from_f32(1.0), 3);
330    test_as_type!(f32, 1.0, bf16, bf16::from_f32(1.0), 3);
331    test_as_type!(f32, 1.0, complex64, complex64::new(1.0, 0.0), 3);
332
333    test_as_type!(f16, f16::from_f32(1.0), bool, true, 3);
334    test_as_type!(f16, f16::from_f32(1.0), i8, 1, 3);
335    test_as_type!(f16, f16::from_f32(1.0), i16, 1, 3);
336    test_as_type!(f16, f16::from_f32(1.0), i32, 1, 3);
337    test_as_type!(f16, f16::from_f32(1.0), i64, 1, 3);
338    test_as_type!(f16, f16::from_f32(1.0), u8, 1, 3);
339    test_as_type!(f16, f16::from_f32(1.0), u16, 1, 3);
340    test_as_type!(f16, f16::from_f32(1.0), u32, 1, 3);
341    test_as_type!(f16, f16::from_f32(1.0), u64, 1, 3);
342    test_as_type!(f16, f16::from_f32(1.0), f32, 1.0, 3);
343    test_as_type!(f16, f16::from_f32(1.0), bf16, bf16::from_f32(1.0), 3);
344    test_as_type!(
345        f16,
346        f16::from_f32(1.0),
347        complex64,
348        complex64::new(1.0, 0.0),
349        3
350    );
351
352    test_as_type!(bf16, bf16::from_f32(1.0), bool, true, 3);
353    test_as_type!(bf16, bf16::from_f32(1.0), i8, 1, 3);
354    test_as_type!(bf16, bf16::from_f32(1.0), i16, 1, 3);
355    test_as_type!(bf16, bf16::from_f32(1.0), i32, 1, 3);
356    test_as_type!(bf16, bf16::from_f32(1.0), i64, 1, 3);
357    test_as_type!(bf16, bf16::from_f32(1.0), u8, 1, 3);
358    test_as_type!(bf16, bf16::from_f32(1.0), u16, 1, 3);
359    test_as_type!(bf16, bf16::from_f32(1.0), u32, 1, 3);
360    test_as_type!(bf16, bf16::from_f32(1.0), u64, 1, 3);
361    test_as_type!(bf16, bf16::from_f32(1.0), f32, 1.0, 3);
362    test_as_type!(bf16, bf16::from_f32(1.0), f16, f16::from_f32(1.0), 3);
363
364    test_as_type!(complex64, complex64::new(1.0, 0.0), bool, true, 3);
365    test_as_type!(complex64, complex64::new(1.0, 0.0), i8, 1, 3);
366    test_as_type!(complex64, complex64::new(1.0, 0.0), i16, 1, 3);
367    test_as_type!(complex64, complex64::new(1.0, 0.0), i32, 1, 3);
368    test_as_type!(complex64, complex64::new(1.0, 0.0), i64, 1, 3);
369    test_as_type!(complex64, complex64::new(1.0, 0.0), u8, 1, 3);
370    test_as_type!(complex64, complex64::new(1.0, 0.0), u16, 1, 3);
371    test_as_type!(complex64, complex64::new(1.0, 0.0), u32, 1, 3);
372    test_as_type!(complex64, complex64::new(1.0, 0.0), u64, 1, 3);
373    test_as_type!(complex64, complex64::new(1.0, 0.0), f32, 1.0, 3);
374    test_as_type!(
375        complex64,
376        complex64::new(1.0, 0.0),
377        f16,
378        f16::from_f32(1.0),
379        3
380    );
381    test_as_type!(
382        complex64,
383        complex64::new(1.0, 0.0),
384        bf16,
385        bf16::from_f32(1.0),
386        3
387    );
388
389    #[test]
390    fn test_view() {
391        let array = Array::from_slice(&[1i16, 2, 3], &[3]);
392        let new_array = array.view::<i8>().unwrap();
393
394        assert_eq!(new_array.dtype(), Dtype::Int8);
395        assert_eq!(new_array.shape(), &[6]);
396        assert_eq!(new_array.item_size(), 1);
397        assert_eq!(new_array.as_slice::<i8>(), &[1, 0, 2, 0, 3, 0]);
398    }
399
400    // The tests below are adapted from the C++ unit test `ops_tests.cpp/test fp8 conversion`
401    #[test]
402    fn test_fp8_conversion() {
403        // Test round-trip for float32
404        let input_f32 = Array::from_slice(&[-1.125f32, -1.0, 0.0, 1.0, 1.125, 4.5, 448.0], &[7]);
405        let fp8 = input_f32.to_fp8().unwrap();
406        assert_eq!(fp8.dtype(), Dtype::Uint8);
407        let output_f32 = fp8.from_fp8(Dtype::Float32).unwrap();
408        assert_eq!(output_f32.dtype(), Dtype::Float32);
409        let data: &[f32] = output_f32.as_slice();
410        assert_eq!(data, &[-1.125f32, -1.0, 0.0, 1.0, 1.125, 4.5, 448.0]);
411
412        // Test round-trip for float16
413        let input_f16 = Array::from_slice(
414            &[
415                f16::from_f32(-1.125),
416                f16::from_f32(-1.0),
417                f16::from_f32(0.0),
418                f16::from_f32(1.0),
419                f16::from_f32(1.125),
420                f16::from_f32(4.5),
421                f16::from_f32(448.0),
422            ],
423            &[7],
424        );
425        let fp8 = input_f16.to_fp8().unwrap();
426        let output_f16 = fp8.from_fp8(Dtype::Float16).unwrap();
427        assert_eq!(output_f16.dtype(), Dtype::Float16);
428        let data: &[f16] = output_f16.as_slice();
429        let expected_f16: Vec<f16> = vec![
430            f16::from_f32(-1.125),
431            f16::from_f32(-1.0),
432            f16::from_f32(0.0),
433            f16::from_f32(1.0),
434            f16::from_f32(1.125),
435            f16::from_f32(4.5),
436            f16::from_f32(448.0),
437        ];
438        assert_eq!(data, expected_f16.as_slice());
439
440        // Test round-trip for bfloat16
441        let input_bf16 = Array::from_slice(
442            &[
443                bf16::from_f32(-1.125),
444                bf16::from_f32(-1.0),
445                bf16::from_f32(0.0),
446                bf16::from_f32(1.0),
447                bf16::from_f32(1.125),
448                bf16::from_f32(4.5),
449                bf16::from_f32(448.0),
450            ],
451            &[7],
452        );
453        let fp8 = input_bf16.to_fp8().unwrap();
454        let output_bf16 = fp8.from_fp8(Dtype::Bfloat16).unwrap();
455        assert_eq!(output_bf16.dtype(), Dtype::Bfloat16);
456        let data: &[bf16] = output_bf16.as_slice();
457        let expected_bf16: Vec<bf16> = vec![
458            bf16::from_f32(-1.125),
459            bf16::from_f32(-1.0),
460            bf16::from_f32(0.0),
461            bf16::from_f32(1.0),
462            bf16::from_f32(1.125),
463            bf16::from_f32(4.5),
464            bf16::from_f32(448.0),
465        ];
466        assert_eq!(data, expected_bf16.as_slice());
467
468        // Test rounding - noisy input should round to expected values
469        let noisy_in =
470            Array::from_slice(&[-1.135f32, -1.01, 0.0001, 1.01, 1.135, 4.6, 447.0], &[7]);
471        let expected = Array::from_slice(&[-1.125f32, -1.0, 0.0, 1.0, 1.125, 4.5, 448.0], &[7]);
472        let fp8 = noisy_in.to_fp8().unwrap();
473        let output = fp8.from_fp8(Dtype::Float32).unwrap();
474        let output_data: &[f32] = output.as_slice();
475        let expected_data: &[f32] = expected.as_slice();
476        assert_eq!(output_data, expected_data);
477
478        // Test overflow - values outside representable range get clamped
479        let overflow_in = Array::from_slice(&[-600.0f32, 600.0], &[2]);
480        let fp8 = overflow_in.to_fp8().unwrap();
481        let output = fp8.from_fp8(Dtype::Float32).unwrap();
482        let data: &[f32] = output.as_slice();
483        assert_eq!(data, &[-448.0f32, 448.0]);
484    }
485}
mlx_rs/ops/conversion.rs

mlx_rs/ops/
conversion.rs
