tiny_recursive_rs/data/
mod.rs

1/// Data loading modules for TRM training
2pub mod numpy_dataset;
3
4pub use numpy_dataset::{NumpyDataset, NumpyDataLoader, DatasetMetadata};
5
6use candle_core::{Result, Tensor, Device};
7
8/// Generic data loader trait
9pub trait BatchDataLoader {
10    /// Get next batch of (input, target) tensors
11    fn next_batch(&mut self, device: &Device) -> Result<Option<(Tensor, Tensor)>>;
12
13    /// Reset loader for new epoch
14    fn reset(&mut self);
15
16    /// Get total number of batches
17    fn num_batches(&self) -> usize;
18}