iris/core/
mat.rs

1use burn::tensor::{Tensor, backend::Backend};
2
3/// A multi-dimensional dense array representation.
4/// It wraps a Burn Tensor and exposes convenient operations.
5#[derive(Clone, Debug)]
6pub struct Mat<B: Backend, const D: usize> {
7    /// The underlying Burn Tensor.
8    pub tensor: Tensor<B, D>,
9}
10
11impl<B: Backend, const D: usize> Mat<B, D> {
12    /// Creates a new Mat from a Burn Tensor.
13    pub fn new(tensor: Tensor<B, D>) -> Self {
14        Self { tensor }
15    }
16
17    /// Returns the shape/dimensions of the Mat.
18    pub fn shape(&self) -> Vec<usize> {
19        self.tensor.dims().to_vec()
20    }
21
22    /// Returns the total number of elements in the Mat.
23    pub fn total(&self) -> usize {
24        self.tensor.shape().num_elements()
25    }
26}
27
28#[cfg(test)]
29mod tests {
30    use super::*;
31    use crate::test_helpers::{TestBackend, test_device};
32    use burn::tensor::{Tensor, TensorData};
33
34    #[test]
35    fn test_mat_operations() {
36        let device = test_device();
37        let tensor = Tensor::<TestBackend, 2>::from_data(
38            TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0], [2, 2]),
39            &device,
40        );
41        let mat = Mat::new(tensor);
42        assert_eq!(mat.shape(), vec![2, 2]);
43        assert_eq!(mat.total(), 4);
44    }
45}
iris/core/mat.rs

iris/core/
mat.rs
