burn_dataset/transform/
composed.rs

1use crate::Dataset;
2use std::error::Error;
3
4/// Compose multiple datasets together to create a bigger one.
5#[derive(new)]
6pub struct ComposedDataset<D> {
7    datasets: Vec<D>,
8}
9
10impl<D, I, E> Dataset<I, E> for ComposedDataset<D>
11where
12    D: Dataset<I, E>,
13    I: Clone,
14    E: Error + Send + Sync + 'static,
15{
16    fn get(&self, index: usize) -> Result<I, E> {
17        let mut current_index = 0;
18        for dataset in self.datasets.iter() {
19            if index < dataset.len() + current_index {
20                return dataset.get(index - current_index);
21            }
22            current_index += dataset.len();
23        }
24        panic!(
25            "Index out of bounds for ComposedDataset: {index} >= {}",
26            self.len()
27        );
28    }
29    fn len(&self) -> usize {
30        let mut total = 0;
31        for dataset in self.datasets.iter() {
32            total += dataset.len();
33        }
34        total
35    }
36}
37
38#[cfg(test)]
39mod tests {
40    use super::*;
41    use crate::FakeDataset;
42
43    #[test]
44    fn test_composed_dataset() {
45        let dataset1 = FakeDataset::<String>::new(10);
46        let dataset2 = FakeDataset::<String>::new(5);
47
48        let items1 = dataset1.iter().map(Result::unwrap).collect::<Vec<_>>();
49        let items2 = dataset2.iter().map(Result::unwrap).collect::<Vec<_>>();
50
51        let composed = ComposedDataset::new(vec![dataset1, dataset2]);
52
53        assert_eq!(composed.len(), 15);
54
55        let expected_items: Vec<String> = items1.iter().chain(items2.iter()).cloned().collect();
56
57        let items = composed.iter().map(Result::unwrap).collect::<Vec<_>>();
58
59        assert_eq!(items, expected_items);
60    }
61}
burn_dataset/transform/composed.rs

burn_dataset/transform/
composed.rs
