inferox_core/
tensor.rs

1use crate::{Backend, DataType, NumericType};
2
3pub trait Tensor: Clone + Send + Sync {
4    type Dtype: DataType;
5    type Backend: Backend;
6
7    fn shape(&self) -> &[usize];
8
9    fn dtype(&self) -> Self::Dtype;
10
11    fn device(&self) -> <Self::Backend as Backend>::Device;
12
13    fn to_device(
14        &self,
15        device: &<Self::Backend as Backend>::Device,
16    ) -> Result<Self, <Self::Backend as Backend>::Error>;
17
18    fn numel(&self) -> usize {
19        self.shape().iter().product()
20    }
21
22    fn reshape(&self, shape: &[usize]) -> Result<Self, <Self::Backend as Backend>::Error>;
23
24    fn contiguous(&self) -> Result<Self, <Self::Backend as Backend>::Error>;
25}
26
27pub trait TensorBuilder<B: Backend + ?Sized> {
28    fn build_from_slice<T: NumericType>(
29        &self,
30        data: &[T],
31        shape: &[usize],
32    ) -> Result<B::Tensor, B::Error>;
33
34    fn build_from_vec<T: NumericType>(
35        &self,
36        data: Vec<T>,
37        shape: &[usize],
38    ) -> Result<B::Tensor, B::Error>;
39
40    fn zeros(&self, shape: &[usize], dtype: impl DataType) -> Result<B::Tensor, B::Error>;
41
42    fn ones(&self, shape: &[usize], dtype: impl DataType) -> Result<B::Tensor, B::Error>;
43
44    fn randn(&self, shape: &[usize], dtype: impl DataType) -> Result<B::Tensor, B::Error>;
45
46    fn with_device(self, device: B::Device) -> Self;
47}