acyclib 0.3.0

ML library for directed acyclic tensor graphs.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

use std::sync::Arc;

use crate::device::{Device, DeviceBuffer, OperationError};

use super::{dense::DenseMatrix, sparse::SparseMatrix};

#[derive(Debug)]
pub enum Matrix<D: Device> {
    Dense(DenseMatrix<D>),
    Sparse(SparseMatrix<D>),
}

impl<D: Device> Matrix<D> {
    pub fn single_size(&self) -> usize {
        match self {
            Self::Dense(dense) => dense.single_size(),
            Self::Sparse(sparse) => sparse.single_size(),
        }
    }

    pub fn batch_size(&self) -> Option<usize> {
        match self {
            Self::Dense(dense) => dense.batch_size(),
            Self::Sparse(sparse) => sparse.batch_size(),
        }
    }

    pub fn size(&self) -> usize {
        match self {
            Self::Dense(dense) => dense.size(),
            Self::Sparse(sparse) => sparse.size(),
        }
    }

    pub fn device(&self) -> Arc<D> {
        match self {
            Self::Dense(dense) => dense.buf.device(),
            Self::Sparse(sparse) => sparse.buf.device(),
        }
    }

    pub fn dense(&self) -> Result<&DenseMatrix<D>, OperationError<D::DeviceError>> {
        if let Self::Dense(matrix) = self { Ok(matrix) } else { Err(OperationError::InvalidTensorFormat) }
    }

    pub fn dense_mut(&mut self) -> Result<&mut DenseMatrix<D>, OperationError<D::DeviceError>> {
        if let Self::Dense(matrix) = self { Ok(matrix) } else { Err(OperationError::InvalidTensorFormat) }
    }

    pub fn sparse(&self) -> Result<&SparseMatrix<D>, OperationError<D::DeviceError>> {
        if let Self::Sparse(matrix) = self { Ok(matrix) } else { Err(OperationError::InvalidTensorFormat) }
    }

    pub fn sparse_mut(&mut self) -> Result<&mut SparseMatrix<D>, OperationError<D::DeviceError>> {
        if let Self::Sparse(matrix) = self { Ok(matrix) } else { Err(OperationError::InvalidTensorFormat) }
    }

    pub fn swap_with(&mut self, other: &mut Self) -> Result<(), OperationError<D::DeviceError>> {
        match other {
            Self::Dense(x) => x.swap_with(self.dense_mut()?)?,
            Self::Sparse(x) => x.swap_with(self.sparse_mut()?)?,
        }

        Ok(())
    }

    pub fn copy_into(&self, other: &mut Self) -> Result<(), OperationError<D::DeviceError>> {
        match other {
            Self::Dense(x) => x.copy_from(self.dense()?)?,
            Self::Sparse(x) => x.copy_from(self.sparse()?)?,
        }

        Ok(())
    }
}