tensor4all_tensorbackend/
tensor_element.rs

1use anyhow::{anyhow, ensure, Result};
2use num_complex::{Complex32, Complex64};
3use tenferro::{DType, Tensor as NativeTensor, TensorScalar};
4
5/// Public scalar element types supported by tensor4all dense/diag constructors.
6///
7/// Implemented for `f32`, `f64`, `Complex32`, and `Complex64`.
8///
9/// # Examples
10///
11/// ```
12/// use tensor4all_tensorbackend::TensorElement;
13///
14/// let t = f64::dense_native_tensor_from_col_major(&[1.0, 2.0], &[2]).unwrap();
15/// assert_eq!(t.shape(), &[2]);
16///
17/// let vals = f64::dense_values_from_native_col_major(&t).unwrap();
18/// assert_eq!(vals, vec![1.0, 2.0]);
19/// ```
20pub trait TensorElement: TensorScalar + Copy + Send + Sync + 'static {
21    /// Build a dense native tensor from column-major data.
22    /// # Errors
23    ///
24    /// Returns an error when the data length does not match the dimensions (a
25    /// /// shape mismatch) or the backend conversion fails.
26    ///
27    fn dense_native_tensor_from_col_major(data: &[Self], dims: &[usize]) -> Result<NativeTensor>;
28
29    /// Build a dense native tensor by moving an owned column-major payload.
30    ///
31    /// This is the allocation-preserving counterpart of
32    /// [`Self::dense_native_tensor_from_col_major`]. The payload must contain
33    /// exactly the product of `dims` elements.
34    ///
35    /// # Errors
36    ///
37    /// Returns an error for a shape mismatch when the data length does not
38    /// equal the dimension product, or for a backend conversion failure.
39    ///
40    /// # Arguments
41    ///
42    /// * `data` - Owned column-major values whose length equals the product of
43    ///   `dims`.
44    /// * `dims` - Logical tensor dimensions in column-major axis order.
45    ///
46    /// # Returns
47    ///
48    /// A native dense tensor that owns the supplied payload without cloning it.
49    ///
50    /// # Examples
51    ///
52    /// ```
53    /// use tensor4all_tensorbackend::TensorElement;
54    ///
55    /// let tensor = f64::dense_native_tensor_from_col_major_owned(
56    ///     vec![1.0, 2.0, 3.0, 4.0],
57    ///     &[2, 2],
58    /// )
59    /// .unwrap();
60    /// assert_eq!(tensor.shape(), &[2, 2]);
61    /// assert_eq!(tensor.as_slice::<f64>().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
62    /// ```
63    fn dense_native_tensor_from_col_major_owned(
64        data: Vec<Self>,
65        dims: &[usize],
66    ) -> Result<NativeTensor>;
67
68    /// Build a diagonal native tensor from column-major diagonal payload data.
69    /// # Errors
70    ///
71    /// Returns an error when the diagonal payload is incompatible (a shape
72    /// /// mismatch) or the backend conversion fails.
73    ///
74    fn diag_native_tensor_from_col_major(
75        data: &[Self],
76        logical_rank: usize,
77    ) -> Result<NativeTensor>;
78
79    /// Build a rank-0 native tensor.
80    /// # Errors
81    ///
82    /// Returns an error when the scalar cannot be converted (a dtype mismatch or
83    /// /// backend failure).
84    ///
85    fn scalar_native_tensor(value: Self) -> Result<NativeTensor>;
86
87    /// Materialize dense column-major values from a native tensor.
88    /// # Errors
89    ///
90    /// Returns an error when the native tensor cannot be materialized (a dtype
91    /// /// mismatch or backend failure).
92    ///
93    fn dense_values_from_native_col_major(tensor: &NativeTensor) -> Result<Vec<Self>>;
94
95    /// Materialize diagonal values from a dense native tensor.
96    /// # Errors
97    ///
98    /// Returns an error when the native tensor cannot be materialized (a dtype
99    /// /// mismatch or backend failure).
100    ///
101    fn diag_values_from_native_temp(tensor: &NativeTensor) -> Result<Vec<Self>>;
102}
103
104fn tensor_dtype_name(dtype: DType) -> &'static str {
105    match dtype {
106        DType::F32 => "f32",
107        DType::F64 => "f64",
108        DType::I32 => "i32",
109        DType::I64 => "i64",
110        DType::Bool => "bool",
111        DType::C32 => "c32",
112        DType::C64 => "c64",
113    }
114}
115
116fn checked_product(dims: &[usize]) -> Result<usize> {
117    dims.iter().try_fold(1usize, |acc, &dim| {
118        acc.checked_mul(dim)
119            .ok_or_else(|| anyhow::anyhow!("dimension product overflow"))
120    })
121}
122
123fn dense_diagonal_values<T: Copy + Default>(diag: &[T], logical_rank: usize) -> Result<Vec<T>> {
124    ensure!(
125        logical_rank >= 1,
126        "diagonal tensor construction requires at least one logical axis"
127    );
128    let diag_len = diag.len();
129    let dims = vec![diag_len; logical_rank];
130    let total_len = checked_product(&dims)?;
131    let mut dense = vec![T::default(); total_len];
132    let diagonal_stride = (0..logical_rank)
133        .scan(1usize, |stride, _| {
134            let current = *stride;
135            *stride = stride.saturating_mul(diag_len);
136            Some(current)
137        })
138        .sum::<usize>();
139    for (i, value) in diag.iter().copied().enumerate() {
140        dense[i * diagonal_stride] = value;
141    }
142    Ok(dense)
143}
144
145macro_rules! impl_tensor_element {
146    ($ty:ty, $dtype:expr) => {
147        impl TensorElement for $ty {
148            fn dense_native_tensor_from_col_major(
149                data: &[Self],
150                dims: &[usize],
151            ) -> Result<NativeTensor> {
152                let expected_len: usize = checked_product(dims)?;
153                ensure!(
154                    data.len() == expected_len,
155                    "dense tensor len {} does not match dims {:?} (expected {})",
156                    data.len(),
157                    dims,
158                    expected_len
159                );
160                Ok(NativeTensor::from_vec_col_major(
161                    dims.to_vec(),
162                    data.to_vec(),
163                )?)
164            }
165
166            fn dense_native_tensor_from_col_major_owned(
167                data: Vec<Self>,
168                dims: &[usize],
169            ) -> Result<NativeTensor> {
170                let expected_len: usize = checked_product(dims)?;
171                ensure!(
172                    data.len() == expected_len,
173                    "dense tensor len {} does not match dims {:?} (expected {})",
174                    data.len(),
175                    dims,
176                    expected_len
177                );
178                Ok(NativeTensor::from_vec_col_major(dims.to_vec(), data)?)
179            }
180
181            fn diag_native_tensor_from_col_major(
182                data: &[Self],
183                logical_rank: usize,
184            ) -> Result<NativeTensor> {
185                let dims = vec![data.len(); logical_rank];
186                let dense = dense_diagonal_values(data, logical_rank)?;
187                Self::dense_native_tensor_from_col_major(&dense, &dims)
188            }
189
190            fn scalar_native_tensor(value: Self) -> Result<NativeTensor> {
191                Ok(NativeTensor::from_vec_col_major(vec![], vec![value])?)
192            }
193
194            fn dense_values_from_native_col_major(tensor: &NativeTensor) -> Result<Vec<Self>> {
195                tensor
196                    .as_slice::<Self>()
197                    .map(|values| values.to_vec())
198                    .map_err(|_| {
199                        anyhow!(
200                            "tensor dtype mismatch: expected {}, got {}",
201                            tensor_dtype_name($dtype),
202                            tensor_dtype_name(tensor.dtype())
203                        )
204                    })
205            }
206
207            fn diag_values_from_native_temp(tensor: &NativeTensor) -> Result<Vec<Self>> {
208                let shape = tensor.shape();
209                ensure!(
210                    !shape.is_empty(),
211                    "diagonal extraction requires rank >= 1, got scalar tensor"
212                );
213                let diag_len = shape[0];
214                ensure!(
215                    shape.iter().all(|&dim| dim == diag_len),
216                    "expected square/equal dims for diagonal extraction, got {:?}",
217                    shape
218                );
219                let dense = Self::dense_values_from_native_col_major(tensor)?;
220                let diagonal_stride = (0..shape.len())
221                    .scan(1usize, |stride, _| {
222                        let current = *stride;
223                        *stride = stride.saturating_mul(diag_len);
224                        Some(current)
225                    })
226                    .sum::<usize>();
227                Ok((0..diag_len).map(|i| dense[i * diagonal_stride]).collect())
228            }
229        }
230    };
231}
232
233impl_tensor_element!(f32, DType::F32);
234impl_tensor_element!(f64, DType::F64);
235impl_tensor_element!(Complex32, DType::C32);
236impl_tensor_element!(Complex64, DType::C64);
tensor4all_tensorbackend/tensor_element.rs

tensor4all_tensorbackend/
tensor_element.rs
