nuvec_core/
creation.rs

1//! Creating arrays from host data: the [`array!`](crate::array!) macro, `From`/`FromIterator`
2//! conversions and [`Array::from_slice`].
3
4use std::panic::Location;
5
6use crate::axes::{self, IntoShape};
7use crate::buffer::Buffer;
8use crate::shape;
9use crate::{Array, Complex64, DType, Element, Error, Result, bf16, f16};
10
11/// An array of `data` in `shape` (row-major).
12pub(crate) fn from_data<T: Element>(data: &[T], shape: &[usize]) -> Result<Array> {
13    let expected = shape::checked_size(shape)?;
14    if data.len() != expected {
15        return Err(Error::size_mismatch("from_slice", expected, data.len()));
16    }
17    let mut buffer = Buffer::try_alloc(shape::checked_nbytes(shape, T::DTYPE.size())?)?;
18    T::write_bytes(data, buffer.bytes_mut());
19    Array::from_buffer(buffer, T::DTYPE, shape)
20}
21
22impl Array {
23    /// Runs the constructor `f` of the public operation `op` on the requested shape
24    /// (`spec`, validated), returning its array; on error, a failed array of `dtype` with
25    /// the requested shape if it is valid, else 0-D.
26    pub(crate) fn constructor(
27        op: &'static str,
28        at: &'static Location<'static>,
29        dtype: DType,
30        spec: &[i64],
31        f: impl FnOnce(&[usize]) -> Result<Array>,
32    ) -> Array {
33        match axes::dims(spec) {
34            Ok(shape) => {
35                f(&shape).unwrap_or_else(|e| Array::failed(e.locate(op, at), dtype, &shape))
36            }
37            Err(e) => Array::failed(e.locate(op, at), dtype, &[]),
38        }
39    }
40
41    /// An array holding `data` (row-major) in `shape`, of `T`'s dtype.
42    ///
43    /// A length that does not match the shape gives a failed array
44    /// ([`crate::OpErrorKind::SizeMismatch`]) with the requested shape.
45    ///
46    /// ```
47    /// use nuvec_core::{Array, DType};
48    ///
49    /// let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]);
50    /// assert_eq!((a.shape(), a.dtype()), (&[2, 3][..], DType::F32));
51    /// assert!(Array::from_slice(&[1u8, 2], [3]).is_failed());
52    /// ```
53    #[track_caller]
54    pub fn from_slice<T: Element>(data: &[T], shape: impl IntoShape) -> Array {
55        let spec = shape.shape_spec();
56        Self::constructor("from_slice", Location::caller(), T::DTYPE, &spec, |shape| {
57            from_data(data, shape)
58        })
59    }
60}
61
62/// A 0-D array holding `value`, of `T`'s dtype (a *strong* scalar: unlike a Rust number used
63/// as an operand, it does not adapt to the other operand's dtype).
64impl<T: Element> From<T> for Array {
65    #[track_caller]
66    fn from(value: T) -> Self {
67        Array::from_slice(&[value], ())
68    }
69}
70
71/// A 1-D array of the elements, of `T`'s dtype.
72impl<T: Element, const N: usize> From<[T; N]> for Array {
73    #[track_caller]
74    fn from(values: [T; N]) -> Self {
75        Array::from_slice(&values, [N])
76    }
77}
78
79/// A 1-D array of the elements, of `T`'s dtype.
80impl<T: Element> From<Vec<T>> for Array {
81    #[track_caller]
82    fn from(values: Vec<T>) -> Self {
83        Array::from_slice(&values, [values.len()])
84    }
85}
86
87/// A 1-D array of the elements, of `T`'s dtype.
88impl<T: Element> From<&[T]> for Array {
89    #[track_caller]
90    fn from(values: &[T]) -> Self {
91        Array::from_slice(values, [values.len()])
92    }
93}
94
95/// A 1-D array of the elements, of `T`'s dtype.
96///
97/// ```
98/// use nuvec_core::{Array, DType};
99///
100/// let a: Array = (0..4).map(|i| i as f32 * 0.5).collect();
101/// assert_eq!((a.shape(), a.dtype()), (&[4][..], DType::F32));
102/// ```
103impl<T: Element> FromIterator<T> for Array {
104    #[track_caller]
105    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
106        let values: Vec<T> = iter.into_iter().collect();
107        Array::from_slice(&values, [values.len()])
108    }
109}
110
111/// An element type of an [`array!`](crate::array!) literal: Rust's default float literal
112/// type `f64` becomes float32 (MLX's default); every other element type keeps its
113/// dtype.
114#[doc(hidden)]
115pub trait Literal: Copy {
116    fn into_array(data: &[Self], shape: &[usize], at: &'static Location<'static>) -> Array;
117}
118
119macro_rules! literal {
120    ($($t:ty),*) => {$(
121        impl Literal for $t {
122            fn into_array(data: &[Self], shape: &[usize], at: &'static Location<'static>) -> Array {
123                from_data(data, shape).unwrap_or_else(|e| {
124                    Array::failed(e.locate("array", at), <$t as Element>::DTYPE, shape)
125                })
126            }
127        }
128    )*};
129}
130
131literal!(
132    bool, u8, u16, u32, u64, i8, i16, i32, i64, f16, bf16, f32, Complex64
133);
134
135impl Literal for f64 {
136    fn into_array(data: &[Self], shape: &[usize], at: &'static Location<'static>) -> Array {
137        let data: Vec<f32> = data.iter().map(|&x| x as f32).collect();
138        f32::into_array(&data, shape, at)
139    }
140}
141
142/// Support of the [`array!`](crate::array!) macro.
143#[doc(hidden)]
144pub mod __private {
145    use std::panic::Location;
146
147    use super::Literal;
148    use crate::Array;
149
150    #[track_caller]
151    pub fn from_nested1<T: Literal, const A: usize>(x: [T; A]) -> Array {
152        T::into_array(&x, &[A], Location::caller())
153    }
154
155    #[track_caller]
156    pub fn from_nested2<T: Literal, const A: usize, const B: usize>(x: [[T; B]; A]) -> Array {
157        let flat: Vec<T> = x.iter().flatten().copied().collect();
158        T::into_array(&flat, &[A, B], Location::caller())
159    }
160
161    #[track_caller]
162    pub fn from_nested3<T: Literal, const A: usize, const B: usize, const C: usize>(
163        x: [[[T; C]; B]; A],
164    ) -> Array {
165        let flat: Vec<T> = x.iter().flatten().flatten().copied().collect();
166        T::into_array(&flat, &[A, B, C], Location::caller())
167    }
168
169    #[track_caller]
170    pub fn from_nested4<
171        T: Literal,
172        const A: usize,
173        const B: usize,
174        const C: usize,
175        const D: usize,
176    >(
177        x: [[[[T; D]; C]; B]; A],
178    ) -> Array {
179        let flat: Vec<T> = x.iter().flatten().flatten().flatten().copied().collect();
180        T::into_array(&flat, &[A, B, C, D], Location::caller())
181    }
182}
183
184/// An array from nested literals, as MLX's `mx.array([...])`: the macro's brackets are the
185/// outermost dimension, so `array![1, 2, 3]` is 1-D and `array![[1.0, 2.0], [3.0, 4.0]]` is
186/// 2-D (up to 4-D). Unsuffixed float literals give float32 and unsuffixed integers int32;
187/// suffixed literals keep their type (`array![1u8, 2, 3]` is uint8), except `f64`, which
188/// gives float32 like MLX's Python floats (use [`Array::from`] for a float64 array).
189/// Ragged rows do not compile.
190///
191/// A 0-D array is [`Array::from`] a number: a macro cannot tell `array!(2.5)` from
192/// `array![2.5]` (the invocation's delimiters are not part of its input), and the latter
193/// is the 1-element vector, as with `vec!`.
194///
195/// ```
196/// use nuvec_core::{Array, DType, array};
197///
198/// let m = array![[1.0, 2.0], [3.0, 4.0]];
199/// assert_eq!((m.shape(), m.dtype()), (&[2, 2][..], DType::F32));
200/// assert_eq!(array![1, 2, 3].dtype(), DType::I32);
201/// assert_eq!(array![1u8, 2, 3].dtype(), DType::U8);
202/// assert_eq!(array![2.5].shape(), [1]);
203/// assert_eq!(Array::from(2.5f32).ndim(), 0);
204/// ```
205#[macro_export]
206macro_rules! array {
207    ($([$([$([$($x:expr),* $(,)?]),* $(,)?]),* $(,)?]),+ $(,)?) => {
208        $crate::__private::from_nested4([$([$([$([$($x),*]),*]),*]),+])
209    };
210    ($([$([$($x:expr),* $(,)?]),* $(,)?]),+ $(,)?) => {
211        $crate::__private::from_nested3([$([$([$($x),*]),*]),+])
212    };
213    ($([$($x:expr),* $(,)?]),+ $(,)?) => {
214        $crate::__private::from_nested2([$([$($x),*]),+])
215    };
216    ($($x:expr),+ $(,)?) => {
217        $crate::__private::from_nested1([$($x),+])
218    };
219}
220
221#[cfg(test)]
222mod tests {
223    use crate::{Array, DType};
224
225    #[test]
226    fn literals() {
227        let a = crate::array![[[1, 2], [3, 4]], [[5, 6], [7, 8]]];
228        assert_eq!((a.shape(), a.dtype()), (&[2, 2, 2][..], DType::I32));
229        let b = crate::array![[[[1.5f64]]]];
230        assert_eq!((b.shape(), b.dtype()), (&[1, 1, 1, 1][..], DType::F32));
231        assert_eq!(crate::array![true, false].dtype(), DType::Bool);
232        assert_eq!(Array::from(vec![1i64; 6]).shape(), [6]);
233        assert_eq!(Array::from(3.0f64).dtype(), DType::F64);
234    }
235}
nuvec_core/creation.rs

nuvec_core/
creation.rs
