lumen_core/base/
shape.rs

1use std::{fmt::Display, vec};
2
3use crate::{Error, Result};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct Shape(pub(crate) Vec<usize>);
7
8impl Shape {
9    pub fn scalar() -> Self {
10        Self(vec![])
11    }
12
13    pub fn is_scalar(&self) -> bool {
14        self.0.is_empty() || (self.0.len() == 1 && self.0[0] == 1)
15    } 
16
17    pub fn rank(&self) -> usize {
18        self.0.len()
19    }
20
21    pub fn dims(&self) -> &[usize] {
22        &self.0
23    }
24
25    pub fn into_dims(self) -> Vec<usize> {
26        self.0
27    }
28
29    pub fn dim(&self, dim: impl Dim) -> Result<usize> {
30        let index = dim.to_index(self, "get dim")?;
31        Ok(self.dims()[index])
32    }
33
34    pub fn element_count(&self) -> usize {
35        self.dims().iter().product()
36    }
37
38    pub fn is_contiguous(&self, stride: &[usize]) -> bool {
39        if self.rank() != stride.len() {
40            return false;
41        }
42        let mut acc = 1;
43        for (&stride, &dim) in stride.iter().zip(self.dims().iter()).rev() {
44            if dim > 1 && stride != acc {
45                return false;
46            }
47            acc *= dim;
48        }
49        true
50    }
51
52    pub fn extend(mut self, additional_dims: &[usize]) -> Self {
53        self.0.extend(additional_dims);
54        self
55    }
56
57
58    /// Check whether the two shapes are compatible for broadcast, and if it is the case return the
59    /// broadcasted shape. This is to be used for binary pointwise ops.
60    /// Copy from https://github.com/huggingface/candle
61    pub fn broadcast_shape_binary_op(&self, rhs: &Self, op: &'static str) -> Result<Shape> {
62        let lhs = self;
63        let lhs_dims = lhs.dims();
64        let rhs_dims = rhs.dims();
65        let lhs_ndims = lhs_dims.len();
66        let rhs_ndims = rhs_dims.len();
67        let bcast_ndims = usize::max(lhs_ndims, rhs_ndims);
68        let mut bcast_dims = vec![0; bcast_ndims];
69        for (idx, bcast_value) in bcast_dims.iter_mut().enumerate() {
70            let rev_idx = bcast_ndims - idx;
71            let l_value = if lhs_ndims < rev_idx {
72                1
73            } else {
74                lhs_dims[lhs_ndims - rev_idx]
75            };
76            let r_value = if rhs_ndims < rev_idx {
77                1
78            } else {
79                rhs_dims[rhs_ndims - rev_idx]
80            };
81            *bcast_value = if l_value == r_value {
82                l_value
83            } else if l_value == 1 {
84                r_value
85            } else if r_value == 1 {
86                l_value
87            } else {
88                Err(Error::ShapeMismatchBinaryOp {
89                    lhs: lhs.clone(),
90                    rhs: rhs.clone(),
91                    op,
92                })?
93            }
94        }
95        Ok(Shape::from(bcast_dims))
96    }
97
98
99    /// Returns an iterator over **dimension coordinates**.
100    ///
101    /// This iterator yields the multi-dimensional coordinates
102    /// (e.g., `[i, j, k, ...]`) of each element in the array, independent
103    /// of the physical storage layout.
104    ///
105    /// Example for shape = (2, 2):
106    /// yields: `[0, 0], [0, 1], [1, 0], [1, 1]`
107    pub fn dim_coordinates(&self) -> DimCoordinates {
108        DimCoordinates::from_shape(self)
109    }
110
111    pub fn dims_coordinates<const N: usize>(&self) -> Result<DimNCoordinates<N>> {
112        DimNCoordinates::<N>::from_shape(self)
113    }
114
115    pub fn dim2_coordinates(&self) -> Result<DimNCoordinates<2>> {
116        DimNCoordinates::<2>::from_shape(self)
117    }
118
119    pub fn dim3_coordinates(&self) -> Result<DimNCoordinates<3>> {
120        DimNCoordinates::<3>::from_shape(self)
121    }
122
123    pub fn dim4_coordinates(&self) -> Result<DimNCoordinates<4>> {
124        DimNCoordinates::<4>::from_shape(self)
125    }
126
127    pub fn dim5_coordinates(&self) -> Result<DimNCoordinates<5>> {
128        DimNCoordinates::<5>::from_shape(self)
129    }
130
131    pub(crate) fn stride_contiguous(&self) -> Vec<usize> {
132        let mut stride = self.dims()
133            .iter()
134            .rev()
135            .scan(1, |prod, u| {
136                let prod_pre_mult = *prod;
137                *prod *= u;
138                Some(prod_pre_mult)
139            })
140            .collect::<Vec<_>>();
141        stride.reverse();
142        stride
143    }
144}
145
146
147//////////////////////////////////////////////////////////////////////////////////////
148///                  DimCoordinates
149//////////////////////////////////////////////////////////////////////////////////////
150
151pub struct DimCoordinates {
152    shape: Vec<usize>,
153    current: Vec<usize>,
154    done: bool,
155}
156
157impl DimCoordinates {
158    pub fn from_shape(shape: &Shape) -> Self {
159        let rank = shape.rank();
160        Self {
161            shape: shape.dims().to_vec(),
162            current: vec![0; rank],
163            done: shape.is_scalar(),
164        }
165    }
166}
167
168impl Iterator for DimCoordinates {
169    type Item = Vec<usize>;
170
171    fn next(&mut self) -> Option<Self::Item> {
172        if self.done {
173            return None;
174        }
175
176        let result = self.current.clone();
177
178        for i in (0..self.current.len()).rev() {
179            self.current[i] += 1;
180            if self.current[i] < self.shape[i] {
181                break; 
182            } else {
183                self.current[i] = 0;
184                if i == 0 {
185                    self.done = true;
186                }
187            }
188        }
189
190        Some(result)
191    }
192}
193
194pub struct DimNCoordinates<const N: usize> {
195    shape: [usize; N],
196    current: [usize; N],
197    done: bool,
198}
199
200impl<const N: usize> DimNCoordinates<N> {
201    pub fn from_shape(from_shape: &Shape) -> Result<Self> {
202        if from_shape.rank() == N {
203            let mut shape = [0usize; N];
204            for i in 0..N {
205                shape[i] = from_shape.dims()[i];
206            }
207
208            let current = [0usize; N];
209            
210            Ok(Self {
211                shape,
212                current,
213                done: N == 0
214            })
215        } else {
216            Err(Error::UnexpectedNumberOfDims {
217                expected: N,
218                got: from_shape.rank(),
219                shape: Shape::from(from_shape.dims()),
220            })?
221        }
222    }
223}
224
225impl<const N: usize> Iterator for DimNCoordinates<N> {
226    type Item = [usize; N];
227    fn next(&mut self) -> Option<Self::Item> {
228        if self.done {
229            return None;
230        }
231
232        let result = self.current;
233
234        for i in (0..N).rev() {
235            self.current[i] += 1;
236            if self.current[i] < self.shape[i] {
237                break; 
238            } else {
239                self.current[i] = 0;
240                if i == 0 {
241                    self.done = true;
242                }
243            }
244        }
245
246        Some(result)
247    }
248}
249
250impl<const C: usize> From<&[usize; C]> for Shape {
251    fn from(dims: &[usize; C]) -> Self {
252        Self(dims.to_vec())
253    }
254}
255
256impl From<Vec<usize>> for Shape {
257    fn from(dims: Vec<usize>) -> Self {
258        Self(dims)
259    }
260}
261
262impl From<&Vec<usize>> for Shape {
263    fn from(dims: &Vec<usize>) -> Self {
264        Self(dims.clone())
265    }
266}
267
268impl From<&[usize]> for Shape {
269    fn from(dims: &[usize]) -> Self {
270        Self(dims.to_vec())
271    }
272}
273
274impl From<&Shape> for Shape {
275    fn from(shape: &Shape) -> Self {
276        Self(shape.0.to_vec())
277    }
278}
279
280impl From<usize> for Shape {
281    fn from(d1: usize) -> Self {
282        Self([d1].to_vec())
283    }
284}
285
286impl From<()> for Shape {
287    fn from(_: ()) -> Self {
288        Self(vec![])
289    }
290}
291
292impl std::fmt::Display for Shape {
293    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
294        write!(f, "(")?;
295        for (i, dim) in self.0.iter().enumerate() {
296            if i > 0 {
297                write!(f, ", ")?;
298            }
299            write!(f, "{}", dim)?;
300        }
301        if self.0.len() == 1 {
302            write!(f, ",")?;
303        }
304        write!(f, ")")
305    }
306}
307
308macro_rules! impl_from_tuple {
309    ($tuple:ty, $($index:tt),+) => {
310        impl From<$tuple> for Shape {
311            fn from(d: $tuple) -> Self {
312                Self([$(d.$index,)+].to_vec())
313            }
314        }
315    };
316}
317
318impl_from_tuple!((usize,), 0);
319impl_from_tuple!((usize, usize), 0, 1);
320impl_from_tuple!((usize, usize, usize), 0, 1, 2);
321impl_from_tuple!((usize, usize, usize, usize), 0, 1, 2, 3);
322impl_from_tuple!((usize, usize, usize, usize, usize), 0, 1, 2, 3, 4);
323impl_from_tuple!((usize, usize, usize, usize, usize, usize), 0, 1, 2, 3, 4, 5);
324
325#[derive(Debug, Clone, Copy, PartialEq, Eq)]
326pub enum D {
327    Minus1,
328    Minus2,
329    Minus(usize),
330    Index(usize),
331}
332
333impl D {
334    pub fn to_real_index(&self, size: usize, op: &'static str) -> Result<usize> {
335        match self {
336            Self::Minus1 if size >= 1 => Ok(size - 1),
337            Self::Minus2 if size >= 2 => Ok(size - 2),
338            Self::Minus(u) if *u > 0 && size >= *u => Ok(size - *u),
339            Self::Index(u) if *u < size => Ok(*u),
340            _ => Err(crate::Error::DimSizeOutOfRange { size, op })?,
341        }
342    }
343}
344
345impl Display for D {
346    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
347        match self {
348            Self::Minus(n) => writeln!(f, "-{}", n),
349            Self::Minus1 => writeln!(f, "-1"),
350            Self::Minus2 => writeln!(f, "-2"),
351            Self::Index(n) => writeln!(f, "{}", n),
352        }
353    }
354}
355
356impl D {
357    fn out_of_range(&self, shape: &Shape, op: &'static str) -> Error {
358        let dim = match self {
359            Self::Minus1 => -1,
360            Self::Minus2 => -2,
361            Self::Minus(u) => -(*u as i32),
362            Self::Index(u) => *u as i32,
363        };
364        Error::DimOutOfRange {
365            shape: shape.clone(),
366            dim,
367            op,
368        }
369    }
370}
371
372
373macro_rules! extract_dims {
374    ($fn_name:ident, $cnt:tt, $dims:expr, $out_type:ty) => {
375        pub fn $fn_name(dims: &[usize]) -> Result<$out_type> {
376            if dims.len() != $cnt {
377                Err(Error::UnexpectedNumberOfDims {
378                    expected: $cnt,
379                    got: dims.len(),
380                    shape: Shape::from(dims),
381                })?
382            } else {
383                Ok($dims(dims))
384            }
385        }
386
387        impl Shape {
388            pub fn $fn_name(&self) -> Result<$out_type> {
389                $fn_name(self.0.as_slice())
390            }
391        }
392
393        impl<T: crate::WithDType> crate::Tensor<T> {
394            pub fn $fn_name(&self) -> Result<$out_type> {
395                self.shape().$fn_name()
396            }
397        }
398
399        impl std::convert::TryInto<$out_type> for Shape {
400            type Error = crate::Error;
401            fn try_into(self) -> crate::Result<$out_type> {
402                self.$fn_name()
403            }
404        }
405    };
406}
407
408extract_dims!(dims0, 0, |_: &[usize]| (), ());
409extract_dims!(dims1, 1, |d: &[usize]| d[0], usize);
410extract_dims!(dims2, 2, |d: &[usize]| (d[0], d[1]), (usize, usize));
411extract_dims!(
412    dims3,
413    3,
414    |d: &[usize]| (d[0], d[1], d[2]),
415    (usize, usize, usize)
416);
417extract_dims!(
418    dims4,
419    4,
420    |d: &[usize]| (d[0], d[1], d[2], d[3]),
421    (usize, usize, usize, usize)
422);
423extract_dims!(
424    dims5,
425    5,
426    |d: &[usize]| (d[0], d[1], d[2], d[3], d[4]),
427    (usize, usize, usize, usize, usize)
428);
429
430
431pub trait Dim : Copy {
432    fn to_index(&self, shape: &Shape, op: &'static str) -> Result<usize>;
433    fn to_index_plus_one(&self, shape: &Shape, op: &'static str) -> Result<usize>;
434}
435
436impl Dim for usize {
437    fn to_index(&self, shape: &Shape, op: &'static str) -> Result<usize> {
438        let dim = *self;
439        if dim >= shape.rank() {
440            Err(Error::DimOutOfRange {
441                shape: shape.clone(),
442                dim: dim as i32,
443                op,
444            })?
445        } else {
446            Ok(dim)
447        }
448    }
449
450    fn to_index_plus_one(&self, shape: &Shape, op: &'static str) -> Result<usize> {
451        let dim = *self;
452        if dim > shape.rank() {
453            Err(Error::DimOutOfRange {
454                shape: shape.clone(),
455                dim: dim as i32,
456                op,
457            })?
458        } else {
459            Ok(dim)
460        }
461    }
462}
463
464impl Dim for D {
465    fn to_index(&self, shape: &Shape, op: &'static str) -> Result<usize> {
466        let rank = shape.rank();
467        match self {
468            // TODO: overfl
469            Self::Minus1 if rank >= 1 => Ok(rank - 1),
470            Self::Minus2 if rank >= 2 => Ok(rank - 2),
471            Self::Minus(u) if *u > 0 && rank >= *u => Ok(rank - *u),
472            Self::Index(u) => u.to_index(shape, op),
473            _ => Err(self.out_of_range(shape, op))?,
474        }
475    }
476
477    fn to_index_plus_one(&self, shape: &Shape, op: &'static str) -> Result<usize> {
478        let rank = shape.rank();
479        match self {
480            Self::Minus1 => Ok(rank),
481            Self::Minus2 if rank >= 1 => Ok(rank - 1),
482            Self::Minus(u) if *u > 0 && rank + 1 >= *u => Ok(rank + 1 - *u),
483            Self::Index(u) => u.to_index_plus_one(shape, op),
484            _ => Err(self.out_of_range(shape, op))?,
485        }
486    }
487}
488
489pub trait Dims {
490    fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>>;
491    fn check_indexes(dims: &[usize], shape: &Shape, op: &'static str) -> Result<()> {
492        for (i, &dim) in dims.iter().enumerate() {
493            if dims[..i].contains(&dim) {
494                return Err(Error::DuplicateDimIndex {
495                    shape: shape.clone(),
496                    dims: dims.to_vec(),
497                    op,
498                })?;
499            }
500            if dim >= shape.rank() {
501                return Err(Error::DimOutOfRange {
502                    shape: shape.clone(),
503                    dim: dim as i32,
504                    op,
505                })?;
506            }
507        }
508        Ok(())
509    }
510}
511
512impl Dims for Vec<usize> {
513    fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
514        Self::check_indexes(&self, shape, op)?;
515        Ok(self)
516    }
517}
518
519impl<const N: usize> Dims for [usize; N] {
520    fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
521        Self::check_indexes(&self, shape, op)?;
522        Ok(self.to_vec())
523    }
524}
525
526impl Dims for &[usize] {
527    fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
528        Self::check_indexes(&self, shape, op)?;
529        Ok(self.to_vec())
530    }
531}
532
533impl Dims for () {
534    fn to_indexes(self, _: &Shape, _: &'static str) -> Result<Vec<usize>> {
535        Ok(vec![])
536    }
537}
538
539impl<D: Dim + Sized> Dims for D {
540    fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
541        let dim = self.to_index(shape, op)?;
542        Ok([dim].to_vec())
543    }
544}
545
546impl<D: Dim> Dims for (D,) {
547    fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
548        let dim = self.0.to_index(shape, op)?;
549        Ok([dim].to_vec())
550    }
551}
552
553impl<D1: Dim, D2: Dim> Dims for (D1, D2) {
554    fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
555        let d0 = self.0.to_index(shape, op)?;
556        let d1 = self.1.to_index(shape, op)?;
557        Ok([d0, d1].to_vec())
558    }
559}
560
561impl<D1: Dim, D2: Dim, D3: Dim> Dims for (D1, D2, D3) {
562    fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
563        let d0 = self.0.to_index(shape, op)?;
564        let d1 = self.1.to_index(shape, op)?;
565        let d2 = self.2.to_index(shape, op)?;
566        Ok([d0, d1, d2].to_vec())
567    }
568}
569
570impl<D1: Dim, D2: Dim, D3: Dim, D4: Dim> Dims for (D1, D2, D3, D4) {
571    fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
572        let d0 = self.0.to_index(shape, op)?;
573        let d1 = self.1.to_index(shape, op)?;
574        let d2 = self.2.to_index(shape, op)?;
575        let d3 = self.3.to_index(shape, op)?;
576        Ok([d0, d1, d2, d3].to_vec())
577    }
578}
579
580impl<D1: Dim, D2: Dim, D3: Dim, D4: Dim, D5: Dim> Dims for (D1, D2, D3, D4, D5) {
581    fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
582        let d0 = self.0.to_index(shape, op)?;
583        let d1 = self.1.to_index(shape, op)?;
584        let d2 = self.2.to_index(shape, op)?;
585        let d3 = self.3.to_index(shape, op)?;
586        let d4 = self.4.to_index(shape, op)?;
587        Ok([d0, d1, d2, d3, d4].to_vec())
588    }
589}
590
591#[cfg(test)]
592mod tests {
593    use super::*;
594
595    #[test]
596    fn stride() {
597        let shape = Shape::from(());
598        assert_eq!(shape.stride_contiguous(), Vec::<usize>::new());
599        let shape = Shape::from(42);
600        assert_eq!(shape.stride_contiguous(), [1]);
601        let shape = Shape::from((42, 1337));
602        assert_eq!(shape.stride_contiguous(), [1337, 1]);
603        let shape = Shape::from((299, 792, 458));
604        assert_eq!(shape.stride_contiguous(), [458 * 792, 458, 1]);
605    }
606
607    #[test]
608    fn test_from_tuple() {
609        let shape = Shape::from((2,));
610        assert_eq!(shape.dims(), &[2]);
611        let shape = Shape::from((2, 3));
612        assert_eq!(shape.dims(), &[2, 3]);
613        let shape = Shape::from((2, 3, 4));
614        assert_eq!(shape.dims(), &[2, 3, 4]);
615        let shape = Shape::from((2, 3, 4, 5));
616        assert_eq!(shape.dims(), &[2, 3, 4, 5]);
617        let shape = Shape::from((2, 3, 4, 5, 6));
618        assert_eq!(shape.dims(), &[2, 3, 4, 5, 6]);
619        let shape = Shape::from((2, 3, 4, 5, 6, 7));
620        assert_eq!(shape.dims(), &[2, 3, 4, 5, 6, 7]);
621    }
622
623    #[test]
624    fn test_dim_coordinates_2d() {
625        let shape = Shape([2, 2].to_vec());
626        let mut iter = shape.dim_coordinates();
627
628        let expected = [
629            [0, 0].to_vec(),
630            [0, 1].to_vec(),
631            [1, 0].to_vec(),
632            [1, 1].to_vec(),
633        ];
634
635        for e in expected {
636            let idx = iter.next();
637            assert_eq!(idx.unwrap(), e);
638        }
639
640        // Iter should be exhausted
641        assert!(iter.next().is_none());
642    }
643
644    #[test]
645    fn test_dim_coordinates_2d_varied() {
646        let shape = Shape([3, 1].to_vec());
647        let mut iter = shape.dim_coordinates();
648
649        let expected = [
650            [0, 0].to_vec(),
651            [1, 0].to_vec(),
652            [2, 0].to_vec(),
653        ];
654
655        for e in expected {
656            let idx = iter.next();
657            assert_eq!(idx.unwrap(), e);
658        }
659
660        assert!(iter.next().is_none());
661    }
662
663    #[test]
664    fn test_dim_coordinates_3d() {
665        let shape = Shape([2, 2, 2].to_vec());
666        let mut iter = shape.dim_coordinates();
667
668        let mut collected = Vec::new();
669        while let Some(idx) = iter.next() {
670            collected.push(idx);
671        }
672
673        let expected = [
674            [0, 0, 0].to_vec(),
675            [0, 0, 1].to_vec(),
676            [0, 1, 0].to_vec(),
677            [0, 1, 1].to_vec(),
678            [1, 0, 0].to_vec(),
679            [1, 0, 1].to_vec(),
680            [1, 1, 0].to_vec(),
681            [1, 1, 1].to_vec(),
682        ];
683
684        assert_eq!(collected, expected);
685    }
686
687    #[test]
688    fn test_dim_n_coordinates_2d() {
689        let shape = Shape([2, 2].to_vec());
690        let mut iter = shape.dim2_coordinates().unwrap();
691
692        let expected = [
693            [0, 0],
694            [0, 1],
695            [1, 0],
696            [1, 1],
697        ];
698
699        for e in expected {
700            let idx = iter.next();
701            assert_eq!(idx.unwrap(), e);
702        }
703
704        assert!(iter.next().is_none());
705    }
706
707    #[test]
708    fn test_dim_n_coordinates_3d() {
709        let shape = Shape([2, 2, 2].to_vec());
710        let mut iter = shape.dim3_coordinates().unwrap();
711
712        let expected = [
713            [0, 0, 0],
714            [0, 0, 1],
715            [0, 1, 0],
716            [0, 1, 1],
717            [1, 0, 0],
718            [1, 0, 1],
719            [1, 1, 0],
720            [1, 1, 1],
721        ];
722
723        for e in expected {
724            let idx = iter.next();
725            assert_eq!(idx.unwrap(), e);
726        }
727
728        assert!(iter.next().is_none());
729    }
730
731    #[test]
732    fn test_dim_n_coordinates_wrong_dim() {
733        let shape = Shape([2, 2].to_vec());
734
735        // dim3_coordinates should return error
736        assert!(shape.dim3_coordinates().is_err());
737        assert!(shape.dims_coordinates::<3>().is_err());
738    }
739
740    #[test]
741    fn test_dim_n_coordinates_empty_shape() {
742        let shape = Shape(vec![]);
743        let mut iter = shape.dims_coordinates::<0>().unwrap();
744        let result = iter.next();
745        assert_eq!(result, None);
746    }
747}
lumen_core/base/shape.rs

lumen_core/base/
shape.rs
