tiny-recursive-rs 0.1.0

Rust implementation of Tiny Recursive Models for efficient puzzle solving
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

/// Data loading modules for TRM training
pub mod numpy_dataset;

pub use numpy_dataset::{NumpyDataset, NumpyDataLoader, DatasetMetadata};

use candle_core::{Result, Tensor, Device};

/// Generic data loader trait
pub trait BatchDataLoader {
    /// Get next batch of (input, target) tensors
    fn next_batch(&mut self, device: &Device) -> Result<Option<(Tensor, Tensor)>>;

    /// Reset loader for new epoch
    fn reset(&mut self);

    /// Get total number of batches
    fn num_batches(&self) -> usize;
}