Refactor FSRSBatcher initialization to remove device parameter, enhancing code consistency across multiple files.

Author: L-M-Sherlock

src/batch_shuffle.rs

@@ -17,7 +17,7 @@ pub(crate) struct BatchTensorDataset<B: Backend> {
 impl<B: Backend> BatchTensorDataset<B> {
     /// Creates a new shuffled dataset.
     pub fn new(dataset: FSRSDataset, batch_size: usize, device: B::Device) -> Self {
-        let batcher = FSRSBatcher::<B>::new(device.clone());
+        let batcher = FSRSBatcher::<B>::new();
         let dataset = dataset
             .items
             .chunks(batch_size)

src/convertor_tests.rs

@@ -304,7 +304,7 @@ fn test_conversion_works() {
     let mut weighted_fsrs_items = constant_weighted_fsrs_items(fsrs_items);
 
     let device = NdArrayDevice::Cpu;
-    let batcher = FSRSBatcher::<NdArrayAutodiff>::new(device);
+    let batcher = FSRSBatcher::<NdArrayAutodiff>::new();
     let res = batcher.batch(vec![weighted_fsrs_items.pop().unwrap()], &device);
     assert_eq!(res.delta_ts.into_scalar(), 64.0);
     assert_eq!(

src/dataset.rs

@@ -79,12 +79,14 @@ impl FSRSItem {
 
 #[derive(Clone)]
 pub(crate) struct FSRSBatcher<B: Backend> {
-    device: B::Device,
+    _backend: core::marker::PhantomData<B>,
 }
 
 impl<B: Backend> FSRSBatcher<B> {
-    pub const fn new(device: B::Device) -> Self {
-        Self { device }
+    pub const fn new() -> Self {
+        Self {
+            _backend: core::marker::PhantomData,
+        }
     }
 }
 
@@ -98,7 +100,7 @@ pub(crate) struct FSRSBatch<B: Backend> {
 }
 
 impl<B: Backend> Batcher<B, WeightedFSRSItem, FSRSBatch<B>> for FSRSBatcher<B> {
-    fn batch(&self, weighted_items: Vec<WeightedFSRSItem>, _device: &B::Device) -> FSRSBatch<B> {
+    fn batch(&self, weighted_items: Vec<WeightedFSRSItem>, device: &B::Device) -> FSRSBatch<B> {
         let pad_size = weighted_items
             .iter()
             .map(|x| x.item.reviews.len())
@@ -123,7 +125,7 @@ impl<B: Backend> Batcher<B, WeightedFSRSItem, FSRSBatch<B>> for FSRSBatcher<B> {
                             dims: vec![1, pad_size],
                         },
                     ),
-                    &self.device,
+                    device,
                 );
                 let rating = Tensor::<B, 2>::from_data(
                     TensorData::new(
@@ -132,7 +134,7 @@ impl<B: Backend> Batcher<B, WeightedFSRSItem, FSRSBatch<B>> for FSRSBatcher<B> {
                             dims: vec![1, pad_size],
                         },
                     ),
-                    &self.device,
+                    device,
                 );
                 (delta_t, rating)
             })
@@ -142,28 +144,24 @@ impl<B: Backend> Batcher<B, WeightedFSRSItem, FSRSBatch<B>> for FSRSBatcher<B> {
             .iter()
             .map(|weighted_item| {
                 let current = weighted_item.item.current();
-                let delta_t: Tensor<B, 1> =
-                    Tensor::from_floats([current.delta_t as f32], &self.device);
+                let delta_t: Tensor<B, 1> = Tensor::from_floats([current.delta_t as f32], device);
                 let label = match current.rating {
                     1 => 0,
                     _ => 1,
                 };
-                let label: Tensor<B, 1, Int> = Tensor::from_ints([label], &self.device);
-                let weight: Tensor<B, 1> =
-                    Tensor::from_floats([weighted_item.weight], &self.device);
+                let label: Tensor<B, 1, Int> = Tensor::from_ints([label], device);
+                let weight: Tensor<B, 1> = Tensor::from_floats([weighted_item.weight], device);
                 (delta_t, label, weight)
             })
             .multiunzip();
 
-        let t_historys = Tensor::cat(time_histories, 0)
-            .transpose()
-            .to_device(&self.device); // [seq_len, batch_size]
+        let t_historys = Tensor::cat(time_histories, 0).transpose().to_device(device); // [seq_len, batch_size]
         let r_historys = Tensor::cat(rating_histories, 0)
             .transpose()
-            .to_device(&self.device); // [seq_len, batch_size]
-        let delta_ts = Tensor::cat(delta_ts, 0).to_device(&self.device);
-        let labels = Tensor::cat(labels, 0).to_device(&self.device);
-        let weights = Tensor::cat(weights, 0).to_device(&self.device);
+            .to_device(device); // [seq_len, batch_size]
+        let delta_ts = Tensor::cat(delta_ts, 0).to_device(device);
+        let labels = Tensor::cat(labels, 0).to_device(device);
+        let weights = Tensor::cat(weights, 0).to_device(device);
 
         // dbg!(&items[0].t_history);
         // dbg!(&t_historys);
@@ -329,7 +327,7 @@ mod tests {
             }
         );
 
-        let batcher = FSRSBatcher::<Backend>::new(DEVICE);
+        let batcher = FSRSBatcher::<Backend>::new();
         use burn::data::dataloader::DataLoaderBuilder;
         let dataloader = DataLoaderBuilder::new(batcher)
             .batch_size(1)
@@ -347,7 +345,7 @@ mod tests {
 
     #[test]
     fn test_batcher() {
-        let batcher = FSRSBatcher::<Backend>::new(DEVICE);
+        let batcher = FSRSBatcher::<Backend>::new();
         let items = [
             FSRSItem {
                 reviews: [(4, 0), (3, 5)]

src/inference.rs

@@ -389,7 +389,7 @@ impl<B: Backend> FSRS<B> {
         }
         let weighted_items = recency_weighted_fsrs_items(items);
         let device = self.device();
-        let batcher = FSRSBatcher::new(device.clone());
+        let batcher = FSRSBatcher::new();
         let mut all_retrievability = vec![];
         let mut all_labels = vec![];
         let mut all_weights = vec![];
@@ -458,7 +458,7 @@ impl<B: Backend> FSRS<B> {
         }
         let weighted_items = constant_weighted_fsrs_items(items);
         let device = self.device();
-        let batcher = FSRSBatcher::new(device.clone());
+        let batcher = FSRSBatcher::new();
         let mut all_predictions_self = vec![];
         let mut all_predictions_other = vec![];
         let mut all_true_val = vec![];
@@ -521,9 +521,9 @@ fn batch_predict(items: Vec<FSRSItem>, parameters: &[f32]) -> Result<Vec<Predict
     }
     let weighted_items = constant_weighted_fsrs_items(items);
     let device = NdArrayDevice::Cpu;
-    let batcher = FSRSBatcher::new(device);
+    let batcher = FSRSBatcher::new();
 
-    let fsrs = FSRS::<NdArray>::new_with_backend(parameters, device)?;
+    let fsrs = FSRS::new(parameters)?;
     let model = fsrs.model();
     let mut predicted_items = Vec::with_capacity(weighted_items.len());
 
@@ -551,7 +551,6 @@ fn batch_predict(items: Vec<FSRSItem>, parameters: &[f32]) -> Result<Vec<Predict
 /// # Returns
 /// A ModelEvaluation containing log loss and RMSE metrics
 fn evaluate(predicted_items: Vec<PredictedFSRSItem>) -> Result<ModelEvaluation> {
-    use burn::backend::NdArray;
     if predicted_items.is_empty() {
         return Err(FSRSError::NotEnoughData);
     }