ruda_tensor/
primitive.rs

1use crate::{Backend, get_device_settings};
2use ruda_core::tensor::{DType, Shape};
3
4#[derive(Debug, Clone)]
5/// A primitive tensor representation.
6pub enum TensorPrimitive<B: Backend> {
7    /// Float tensor primitive.
8    Float(B::FloatTensorPrimitive),
9    /// Quantized float tensor primitive.
10    QFloat(B::QuantizedTensorPrimitive),
11}
12
13impl<B: Backend> TensorPrimitive<B> {
14    /// Returns the full tensor representation.
15    pub fn tensor(self) -> B::FloatTensorPrimitive {
16        match self {
17            Self::QFloat(tensor) => {
18                let dtype = get_device_settings::<B>(&B::q_device(&tensor)).float_dtype;
19                B::dequantize(tensor, dtype)
20            }
21            Self::Float(tensor) => tensor,
22        }
23    }
24
25    /// Returns a mutable reference to the full tensor representation.
26    pub fn get_mut_ref(&mut self) -> &mut B::FloatTensorPrimitive {
27        match self {
28            Self::QFloat(_tensor) => todo!(),
29            Self::Float(tensor) => tensor,
30        }
31    }
32}
33
34impl<B: Backend> TensorMetadata for TensorPrimitive<B> {
35    fn dtype(&self) -> DType {
36        match self {
37            TensorPrimitive::Float(tensor) => tensor.dtype(),
38            TensorPrimitive::QFloat(tensor) => tensor.dtype(),
39        }
40    }
41
42    fn shape(&self) -> Shape {
43        match self {
44            TensorPrimitive::Float(tensor) => tensor.shape(),
45            TensorPrimitive::QFloat(tensor) => tensor.shape(),
46        }
47    }
48
49    fn rank(&self) -> usize {
50        match self {
51            TensorPrimitive::Float(tensor) => tensor.rank(),
52            TensorPrimitive::QFloat(tensor) => tensor.rank(),
53        }
54    }
55}
56
57pub use ruda_core::tensor::primitive::{QTensorPrimitive, TensorMetadata};
ruda_tensor/primitive.rs

ruda_tensor/
primitive.rs
