1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

/// Implementation for Sequence.
///
/// Currently `BinaryHeap`, `BTreeSet`, `HashSet` and `LinkedList` are not supported because the current implementation
/// require indexing for doing the transpose.
///
use super::super::Collate;
use super::DefaultCollate;
use std::collections::VecDeque;

impl<T> Collate<Vec<T>> for DefaultCollate
where
    Self: Collate<T>,
    T: Clone,
{
    type Output = Vec<<Self as Collate<T>>::Output>;
    fn collate(&self, batch: Vec<Vec<T>>) -> Self::Output {
        let elem_size = batch
            .get(0)
            .expect("Batch should contain at least one element")
            .len();

        assert!(
            batch.iter().all(|vec| vec.len() == elem_size),
            "Each Vec in the batch should have equal size"
        );

        let mut collated = Vec::with_capacity(batch.len());

        for i in 0..batch[0].len() {
            let vec: Vec<_> = batch.iter().map(|sample| sample[i].clone()).collect();
            collated.push(self.collate(vec));
        }
        collated
    }
}

impl<T> Collate<VecDeque<T>> for DefaultCollate
where
    Self: Collate<T>,
    T: Clone,
{
    type Output = Vec<<Self as Collate<T>>::Output>;
    fn collate(&self, batch: Vec<VecDeque<T>>) -> Self::Output {
        let elem_size = batch
            .get(0)
            .expect("Batch should contain at least one element")
            .len();

        assert!(
            batch.iter().all(|vec| vec.len() == elem_size),
            "Each Vec in the batch should have equal size"
        );

        let mut collated = Vec::with_capacity(batch.len());

        for i in 0..batch[0].len() {
            let vec: Vec<_> = batch.iter().map(|sample| sample[i].clone()).collect();
            collated.push(self.collate(vec));
        }
        collated
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::array;

    #[test]
    fn vec_of_vec() {
        assert_eq!(
            DefaultCollate::default().collate(vec![vec![1]]),
            vec![array![1]]
        );
        assert_eq!(
            DefaultCollate::default().collate(vec![vec![1, 2], vec![3, 4]]),
            vec![array![1, 3], array![2, 4]]
        );
        // different type
        assert_eq!(
            DefaultCollate::default().collate(vec![vec![true, false], vec![true, false]]),
            vec![array![true, true], array![false, false]]
        );

        assert_eq!(
            DefaultCollate::default().collate(vec![vec![1, 2, 3], vec![4, 5, 6]]),
            vec![array![1, 4], array![2, 5], array![3, 6]]
        );
        // batch_size 3
        assert_eq!(
            DefaultCollate::default().collate(vec![vec![1, 2], vec![3, 4], vec![5, 6]]),
            vec![array![1, 3, 5], array![2, 4, 6]]
        );
        // batch_size 10
        assert_eq!(
            DefaultCollate::default().collate(vec![
                vec![1, 2],
                vec![3, 4],
                vec![5, 6],
                vec![7, 8],
                vec![9, 10],
                vec![11, 12],
                vec![13, 14],
                vec![15, 16],
                vec![17, 18],
                vec![19, 20]
            ]),
            vec![
                array![1, 3, 5, 7, 9, 11, 13, 15, 17, 19],
                array![2, 4, 6, 8, 10, 12, 14, 16, 18, 20]
            ]
        );
    }

    #[test]
    fn specialized() {
        assert_eq!(
            DefaultCollate::default().collate(vec![
                vec![String::from("a"), String::from("b")],
                vec![String::from("c"), String::from("d")]
            ]),
            vec![
                vec![String::from('a'), String::from('c')],
                vec![String::from('b'), String::from('d')],
            ]
        );
    }
}