mlx_rs/ops/
other.rs

1use std::ffi::CString;
2
3use mlx_internal_macros::generate_macro;
4
5use crate::utils::guard::Guarded;
6use crate::utils::VectorArray;
7use crate::{
8    error::{Exception, Result},
9    Array, Dtype, Stream,
10};
11
12/// Diagonal selection and dtype control for [`Array::trace`].
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub struct TraceOptions {
15    /// Diagonal offset.
16    pub offset: i32,
17
18    /// First matrix axis.
19    pub axis1: i32,
20
21    /// Second matrix axis.
22    pub axis2: i32,
23
24    /// Dtype used before reduction.
25    pub dtype: Option<Dtype>,
26}
27
28impl Default for TraceOptions {
29    fn default() -> Self {
30        Self {
31            offset: 0,
32            axis1: 0,
33            axis2: 1,
34            dtype: None,
35        }
36    }
37}
38
39impl Array {
40    /// Sum a selected diagonal.
41    ///
42    /// The diagonal is cast to the requested dtype before reduction. Reduction promotion can
43    /// therefore produce a different final dtype.
44    ///
45    /// ```rust
46    /// use mlx_rs::{array, ops::TraceOptions, Dtype};
47    ///
48    /// let result = array!([[1.0, 2.0], [3.0, 4.0]])
49    ///     .trace(TraceOptions {
50    ///         dtype: Some(Dtype::Int32),
51    ///         ..TraceOptions::default()
52    ///     })
53    ///     .unwrap();
54    /// assert!(result.shape().is_empty());
55    /// ```
56    pub fn trace(&self, options: TraceOptions) -> Result<Array> {
57        let stream = Stream::thread_local_or_default();
58        if options == TraceOptions::default() {
59            Array::try_from_op(|res| unsafe {
60                mlx_sys::mlx_trace(res, self.as_ptr(), stream.as_ref().as_ptr())
61            })
62        } else {
63            Array::try_from_op(|res| unsafe {
64                mlx_sys::mlx_trace_axes(
65                    res,
66                    self.as_ptr(),
67                    options.offset,
68                    options.axis1,
69                    options.axis2,
70                    options.dtype.unwrap_or_else(|| self.dtype()).into(),
71                    stream.as_ref().as_ptr(),
72                )
73            })
74        }
75    }
76
77    /// Extract a diagonal or construct a diagonal matrix.
78    ///
79    /// If self is 1-D then a diagonal matrix is constructed with self on the `k`-th diagonal. If
80    /// self is 2-D then the `k`-th diagonal is returned.
81    ///
82    /// # Params:
83    ///
84    /// - `k`: the diagonal to extract or construct
85    /// - `stream`: stream or device to evaluate on
86    pub fn diag(&self, k: impl Into<Option<i32>>) -> Result<Array> {
87        let stream = Stream::thread_local_or_default();
88        Array::try_from_op(|res| unsafe {
89            mlx_sys::mlx_diag(
90                res,
91                self.as_ptr(),
92                k.into().unwrap_or(0),
93                stream.as_ref().as_ptr(),
94            )
95        })
96    }
97
98    /// Compatibility shim for [`diag`].
99    #[deprecated(
100        since = "0.26.0",
101        note = "use `with_stream` or `with_device` around `diag`"
102    )]
103    pub fn diag_device(
104        &self,
105        k: impl Into<Option<i32>>,
106        stream: impl AsRef<Stream>,
107    ) -> Result<Array> {
108        crate::with_stream(stream.as_ref(), || self.diag(k))
109    }
110
111    /// Return specified diagonals.
112    ///
113    /// If self is 2-D, then a 1-D array containing the diagonal at the given `offset` is returned.
114    ///
115    /// If self has more than two dimensions, then `axis1` and `axis2` determine the 2D subarrays
116    /// from which diagonals are extracted. The new shape is the original shape with `axis1` and
117    /// `axis2` removed and a new dimension inserted at the end corresponding to the diagonal.
118    ///
119    /// # Params:
120    ///
121    /// - `offset`: offset of the diagonal.  Can be positive or negative
122    /// - `axis1`: first axis of the 2-D sub-array from which the diagonals should be taken
123    /// - `axis2`: second axis of the 2-D sub-array from which the diagonals should be taken
124    /// - `stream`: stream or device to evaluate on
125    pub fn diagonal(
126        &self,
127        offset: impl Into<Option<i32>>,
128        axis1: impl Into<Option<i32>>,
129        axis2: impl Into<Option<i32>>,
130    ) -> Result<Array> {
131        let stream = Stream::thread_local_or_default();
132        Array::try_from_op(|res| unsafe {
133            mlx_sys::mlx_diagonal(
134                res,
135                self.as_ptr(),
136                offset.into().unwrap_or(0),
137                axis1.into().unwrap_or(0),
138                axis2.into().unwrap_or(1),
139                stream.as_ref().as_ptr(),
140            )
141        })
142    }
143
144    /// Compatibility shim for [`diagonal`].
145    #[deprecated(
146        since = "0.26.0",
147        note = "use `with_stream` or `with_device` around `diagonal`"
148    )]
149    pub fn diagonal_device(
150        &self,
151        offset: impl Into<Option<i32>>,
152        axis1: impl Into<Option<i32>>,
153        axis2: impl Into<Option<i32>>,
154        stream: impl AsRef<Stream>,
155    ) -> Result<Array> {
156        crate::with_stream(stream.as_ref(), || self.diagonal(offset, axis1, axis2))
157    }
158
159    /// Perform the Walsh-Hadamard transform along the final axis.
160    ///
161    /// Supports sizes `n = m*2^k` for `m` in `(1, 12, 20, 28)` and `2^k <= 8192`
162    /// for ``DType/float32`` and `2^k <= 16384` for ``DType/float16`` and ``DType/bfloat16``.
163    ///
164    /// # Params
165    /// - scale: scale the output by this factor -- default is `1.0/sqrt(array.dim(-1))`
166    /// - stream: stream to evaluate on.
167    pub fn hadamard_transform(&self, scale: impl Into<Option<f32>>) -> Result<Array> {
168        let stream = Stream::thread_local_or_default();
169        let scale = scale.into();
170        let scale = mlx_sys::mlx_optional_float {
171            value: scale.unwrap_or(0.0),
172            has_value: scale.is_some(),
173        };
174
175        Array::try_from_op(|res| unsafe {
176            mlx_sys::mlx_hadamard_transform(res, self.as_ptr(), scale, stream.as_ref().as_ptr())
177        })
178    }
179
180    /// Compatibility shim for [`hadamard_transform`].
181    #[deprecated(
182        since = "0.26.0",
183        note = "use `with_stream` or `with_device` around `hadamard_transform`"
184    )]
185    pub fn hadamard_transform_device(
186        &self,
187        scale: impl Into<Option<f32>>,
188        stream: impl AsRef<Stream>,
189    ) -> Result<Array> {
190        crate::with_stream(stream.as_ref(), || self.hadamard_transform(scale))
191    }
192}
193
194/// See [`Array::diag`]
195pub fn diag(a: impl AsRef<Array>, k: impl Into<Option<i32>>) -> Result<Array> {
196    a.as_ref().diag(k)
197}
198
199/// Compatibility shim for [`diag`].
200#[generate_macro(customize(forwarding_shim = true))]
201#[deprecated(
202    since = "0.26.0",
203    note = "use `with_stream` or `with_device` around `diag`"
204)]
205pub fn diag_device(
206    a: impl AsRef<Array>,
207    #[optional] k: impl Into<Option<i32>>,
208    #[optional] stream: impl AsRef<Stream>,
209) -> Result<Array> {
210    crate::with_stream(stream.as_ref(), || diag(a, k))
211}
212
213/// See [`Array::diagonal`]
214pub fn diagonal(
215    a: impl AsRef<Array>,
216    offset: impl Into<Option<i32>>,
217    axis1: impl Into<Option<i32>>,
218    axis2: impl Into<Option<i32>>,
219) -> Result<Array> {
220    a.as_ref().diagonal(offset, axis1, axis2)
221}
222
223/// Compatibility shim for [`diagonal`].
224#[generate_macro(customize(forwarding_shim = true))]
225#[deprecated(
226    since = "0.26.0",
227    note = "use `with_stream` or `with_device` around `diagonal`"
228)]
229pub fn diagonal_device(
230    a: impl AsRef<Array>,
231    #[optional] offset: impl Into<Option<i32>>,
232    #[optional] axis1: impl Into<Option<i32>>,
233    #[optional] axis2: impl Into<Option<i32>>,
234    #[optional] stream: impl AsRef<Stream>,
235) -> Result<Array> {
236    crate::with_stream(stream.as_ref(), || diagonal(a, offset, axis1, axis2))
237}
238
239/// Perform the Einstein summation convention on the operands.
240///
241/// # Params
242///
243/// - subscripts: Einstein summation convention equation
244/// - operands: input arrays
245/// - stream: stream or device to evaluate on
246pub fn einsum<'a>(
247    subscripts: &str,
248    operands: impl IntoIterator<Item = &'a Array>,
249) -> Result<Array> {
250    let stream = Stream::thread_local_or_default();
251    let c_subscripts =
252        CString::new(subscripts).map_err(|_| Exception::from("Invalid subscripts"))?;
253    let c_operands = VectorArray::try_from_iter(operands.into_iter())?;
254
255    Array::try_from_op(|res| unsafe {
256        mlx_sys::mlx_einsum(
257            res,
258            c_subscripts.as_ptr(),
259            c_operands.as_ptr(),
260            stream.as_ref().as_ptr(),
261        )
262    })
263}
264
265/// Compatibility shim for [`einsum`].
266#[generate_macro(customize(forwarding_shim = true))]
267#[deprecated(
268    since = "0.26.0",
269    note = "use `with_stream` or `with_device` around `einsum`"
270)]
271pub fn einsum_device<'a>(
272    subscripts: &str,
273    operands: impl IntoIterator<Item = &'a Array>,
274    #[optional] stream: impl AsRef<Stream>,
275) -> Result<Array> {
276    crate::with_stream(stream.as_ref(), || einsum(subscripts, operands))
277}
278
279/// Perform the Kronecker product of two arrays.
280///
281/// # Params
282///
283/// - `a`: first array
284/// - `b`: second array
285/// - `stream`: stream or device to evaluate on
286pub fn kron(a: impl AsRef<Array>, b: impl AsRef<Array>) -> Result<Array> {
287    let stream = Stream::thread_local_or_default();
288    Array::try_from_op(|res| unsafe {
289        mlx_sys::mlx_kron(
290            res,
291            a.as_ref().as_ptr(),
292            b.as_ref().as_ptr(),
293            stream.as_ref().as_ptr(),
294        )
295    })
296}
297
298/// Compatibility shim for [`kron`].
299#[generate_macro(customize(forwarding_shim = true))]
300#[deprecated(
301    since = "0.26.0",
302    note = "use `with_stream` or `with_device` around `kron`"
303)]
304pub fn kron_device(
305    a: impl AsRef<Array>,
306    b: impl AsRef<Array>,
307    #[optional] stream: impl AsRef<Stream>,
308) -> Result<Array> {
309    crate::with_stream(stream.as_ref(), || kron(a, b))
310}
311
312#[cfg(test)]
313mod tests {
314    use crate::{
315        array,
316        ops::{arange, diag, einsum, reshape},
317        test_utils::{assert_array_eq, assert_array_eq_with_context, tolerances},
318        Array,
319    };
320    use pretty_assertions::assert_eq;
321
322    use super::diagonal;
323
324    #[test]
325    fn test_diagonal() {
326        let x = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7], &[4, 2]);
327        let out = diagonal(&x, None, None, None).unwrap();
328        assert_array_eq_with_context(
329            out,
330            array!([0, 3]),
331            tolerances::EXACT.rtol,
332            tolerances::EXACT.atol,
333            "default 2d diagonal",
334        );
335
336        assert!(diagonal(&x, 1, 6, 0).is_err());
337        assert!(diagonal(&x, 1, 0, -3).is_err());
338
339        let x = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], &[3, 4]);
340        let out = diagonal(&x, 2, 1, 0).unwrap();
341        assert_array_eq_with_context(
342            out,
343            array!([8]),
344            tolerances::EXACT.rtol,
345            tolerances::EXACT.atol,
346            "positive offset with swapped axes",
347        );
348
349        let out = diagonal(&x, -1, 0, 1).unwrap();
350        assert_array_eq_with_context(
351            out,
352            array!([4, 9]),
353            tolerances::EXACT.rtol,
354            tolerances::EXACT.atol,
355            "negative offset",
356        );
357
358        let out = diagonal(&x, -5, 0, 1).unwrap();
359        out.eval().unwrap();
360        assert_eq!(out.shape(), &[0]);
361
362        let x = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], &[3, 2, 2]);
363        let out = diagonal(&x, 1, 0, 1).unwrap();
364        assert_array_eq_with_context(
365            out,
366            array!([[2], [3]]),
367            tolerances::EXACT.rtol,
368            tolerances::EXACT.atol,
369            "3d positive offset",
370        );
371
372        let out = diagonal(&x, 0, 2, 0).unwrap();
373        assert_array_eq_with_context(
374            out,
375            array!([[0, 5], [2, 7]]),
376            tolerances::EXACT.rtol,
377            tolerances::EXACT.atol,
378            "3d axes 2 and 0",
379        );
380
381        let out = diagonal(&x, 1, -1, 0).unwrap();
382        assert_array_eq_with_context(
383            out,
384            array!([[4, 9], [6, 11]]),
385            tolerances::EXACT.rtol,
386            tolerances::EXACT.atol,
387            "3d negative axis",
388        );
389
390        let x = reshape(arange::<_, f32>(None, 16, None).unwrap(), &[2, 2, 2, 2]).unwrap();
391        let out = diagonal(&x, 0, 0, 1).unwrap();
392        assert_array_eq_with_context(
393            out,
394            Array::from_slice(&[0.0, 12.0, 1.0, 13.0, 2.0, 14.0, 3.0, 15.0], &[2, 2, 2]),
395            tolerances::EXACT.rtol,
396            tolerances::EXACT.atol,
397            "float32 diagonal axes",
398        );
399
400        assert!(diagonal(&x, 0, 1, 1).is_err());
401
402        let x = array!([0, 1]);
403        assert!(diagonal(&x, 0, 0, 1).is_err());
404    }
405
406    #[test]
407    fn test_diag() {
408        // Too few or too many dimensions
409        assert!(diag(Array::from_f32(0.0), None).is_err());
410        assert!(diag(Array::from_slice(&[0.0], &[1, 1, 1]), None).is_err());
411
412        // Test with 1D array
413        let x = array!([0, 1, 2, 3]);
414        let out = diag(&x, 0).unwrap();
415        assert_array_eq(
416            out,
417            array!([[0, 0, 0, 0], [0, 1, 0, 0], [0, 0, 2, 0], [0, 0, 0, 3]]),
418            tolerances::EXACT.rtol,
419            tolerances::EXACT.atol,
420        );
421
422        let out = diag(&x, 1).unwrap();
423        assert_array_eq(
424            out,
425            array!([
426                [0, 0, 0, 0, 0],
427                [0, 0, 1, 0, 0],
428                [0, 0, 0, 2, 0],
429                [0, 0, 0, 0, 3],
430                [0, 0, 0, 0, 0]
431            ]),
432            tolerances::EXACT.rtol,
433            tolerances::EXACT.atol,
434        );
435
436        let out = diag(&x, -1).unwrap();
437        assert_array_eq(
438            out,
439            array!([
440                [0, 0, 0, 0, 0],
441                [0, 0, 0, 0, 0],
442                [0, 1, 0, 0, 0],
443                [0, 0, 2, 0, 0],
444                [0, 0, 0, 3, 0]
445            ]),
446            tolerances::EXACT.rtol,
447            tolerances::EXACT.atol,
448        );
449
450        // Test with 2D array
451        let x = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8], &[3, 3]);
452        let out = diag(&x, 0).unwrap();
453        assert_array_eq(
454            out,
455            array!([0, 4, 8]),
456            tolerances::EXACT.rtol,
457            tolerances::EXACT.atol,
458        );
459
460        let out = diag(&x, 1).unwrap();
461        assert_array_eq(
462            out,
463            array!([1, 5]),
464            tolerances::EXACT.rtol,
465            tolerances::EXACT.atol,
466        );
467
468        let out = diag(&x, -1).unwrap();
469        assert_array_eq(
470            out,
471            array!([3, 7]),
472            tolerances::EXACT.rtol,
473            tolerances::EXACT.atol,
474        );
475    }
476
477    #[test]
478    fn test_einsum() {
479        // Test dot product (vector-vector)
480        let a = array!([0.0, 1.0, 2.0, 3.0]);
481        let b = array!([4.0, 5.0, 6.0, 7.0]);
482        let out = einsum("i,i->", &[a, b]).unwrap();
483        assert_array_eq_with_context(
484            out,
485            array!(38.0),
486            tolerances::EXACT.rtol,
487            tolerances::EXACT.atol,
488            "float32 vector dot",
489        );
490
491        // Test trace (diagonal sum)
492        let m = array!([[1, 2], [3, 4]]);
493        let out = einsum("ii->", &[m]).unwrap();
494        assert_array_eq_with_context(
495            out,
496            array!(5),
497            tolerances::EXACT.rtol,
498            tolerances::EXACT.atol,
499            "int32 matrix trace",
500        );
501    }
502
503    #[test]
504    fn test_hadamard_transform() {
505        let input = Array::from_slice(&[1.0, -1.0, -1.0, 1.0], &[2, 2]);
506        let expected = Array::from_slice(
507            &[
508                0.0,
509                2.0_f32 / 2.0_f32.sqrt(),
510                0.0,
511                -2.0_f32 / 2.0_f32.sqrt(),
512            ],
513            &[2, 2],
514        );
515        let result = input.hadamard_transform(None).unwrap();
516
517        let c = result.all_close(&expected, 1e-5, 1e-5, None).unwrap();
518        assert!(c);
519    }
520
521    // This test is adapted from the python unit test `mlx/test/test_ops.py` `test_kron`
522    #[test]
523    fn test_kron() {
524        // Basic vector test
525        let x = array!([1, 2]);
526        let y = array!([3, 4]);
527        let z = super::kron(&x, &y).unwrap();
528        assert_array_eq(
529            z,
530            array!([3, 4, 6, 8]),
531            tolerances::EXACT.rtol,
532            tolerances::EXACT.atol,
533        );
534
535        // Basic matrix test
536        let x = array!([[1, 2], [3, 4]]);
537        let y = array!([[0, 5], [6, 7]]);
538        let z = super::kron(&x, &y).unwrap();
539        assert_array_eq(
540            z,
541            array!([
542                [0, 5, 0, 10],
543                [6, 7, 12, 14],
544                [0, 15, 0, 20],
545                [18, 21, 24, 28]
546            ]),
547            tolerances::EXACT.rtol,
548            tolerances::EXACT.atol,
549        );
550    }
551}
mlx_rs/ops/other.rs

mlx_rs/ops/
other.rs
