ruda_model/record/
sparse.rs

1use core::marker::PhantomData;
2use ruda_tensor::{api::sparse::CsrTensor, ops::SparseOps};
3use serde::{Serialize, de::DeserializeOwned};
4use super::{PrecisionSettings, Record};
5
6#[derive(Clone, Debug)]
7pub struct CsrTensorRecord<B: SparseOps> {
8    data: B::CsrData,
9    marker: PhantomData<fn() -> B>,
10}
11
12impl<B: SparseOps> CsrTensorRecord<B> {
13    pub fn from_data(data: B::CsrData) -> Self {
14        Self { data, marker: PhantomData }
15    }
16
17    pub fn into_data(self) -> B::CsrData {
18        self.data
19    }
20
21    pub async fn capture(tensor: &CsrTensor<B>) -> Result<Self, B::SparseError> {
22        tensor.to_data().await.map(Self::from_data)
23    }
24
25    pub fn restore(&self, device: &B::Device) -> Result<CsrTensor<B>, B::SparseError> {
26        CsrTensor::from_data(&self.data, device)
27    }
28}
29
30impl<B: SparseOps> Record<B> for CsrTensorRecord<B>
31where
32    B::CsrData: Serialize + DeserializeOwned,
33{
34    type Item<S: PrecisionSettings> = B::CsrData;
35
36    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
37        self.data
38    }
39
40    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, _device: &B::Device) -> Self {
41        Self::from_data(item)
42    }
43}
ruda_model/record/sparse.rs

ruda_model/record/
sparse.rs
