fix: get_slice should not use intermediate Vec (#894)

docs/python/src/conf.py

@@ -1,7 +1,7 @@
 import ezkl
 
 project = 'ezkl'
-release = '16.2.6'
+release = '0.0.0'
 version = release
 
 

src/graph/utilities.rs

@@ -142,8 +142,6 @@ use tract_onnx::prelude::SymbolValues;
 pub fn extract_tensor_value(
     input: Arc<tract_onnx::prelude::Tensor>,
 ) -> Result<Tensor<f32>, GraphError> {
-    use maybe_rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
-
     let dt = input.datum_type();
     let dims = input.shape().to_vec();
 
@@ -156,7 +154,7 @@ pub fn extract_tensor_value(
     match dt {
         DatumType::F16 => {
             let vec = input.as_slice::<tract_onnx::prelude::f16>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| (*x).into()).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| (*x).into()).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::F32 => {
@@ -165,61 +163,61 @@ pub fn extract_tensor_value(
         }
         DatumType::F64 => {
             let vec = input.as_slice::<f64>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| *x as f32).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| *x as f32).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::I64 => {
             // Generally a shape or hyperparam
             let vec = input.as_slice::<i64>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| *x as f32).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| *x as f32).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::I32 => {
             // Generally a shape or hyperparam
             let vec = input.as_slice::<i32>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| *x as f32).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| *x as f32).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::I16 => {
             // Generally a shape or hyperparam
             let vec = input.as_slice::<i16>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| *x as f32).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| *x as f32).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::I8 => {
             // Generally a shape or hyperparam
             let vec = input.as_slice::<i8>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| *x as f32).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| *x as f32).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::U8 => {
             // Generally a shape or hyperparam
             let vec = input.as_slice::<u8>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| *x as f32).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| *x as f32).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::U16 => {
             // Generally a shape or hyperparam
             let vec = input.as_slice::<u16>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| *x as f32).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| *x as f32).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::U32 => {
             // Generally a shape or hyperparam
             let vec = input.as_slice::<u32>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| *x as f32).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| *x as f32).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::U64 => {
             // Generally a shape or hyperparam
             let vec = input.as_slice::<u64>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| *x as f32).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| *x as f32).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::Bool => {
             // Generally a shape or hyperparam
             let vec = input.as_slice::<bool>()?.to_vec();
-            let cast: Vec<f32> = vec.par_iter().map(|x| *x as usize as f32).collect();
+            let cast: Vec<f32> = vec.iter().map(|x| *x as usize as f32).collect();
             const_value = Tensor::<f32>::new(Some(&cast), &dims)?;
         }
         DatumType::TDim => {
@@ -227,7 +225,7 @@ pub fn extract_tensor_value(
             let vec = input.as_slice::<tract_onnx::prelude::TDim>()?.to_vec();
 
             let cast: Result<Vec<f32>, GraphError> = vec
-                .par_iter()
+                .iter()
                 .map(|x| match x.to_i64() {
                     Ok(v) => Ok(v as f32),
                     Err(_) => match x.to_i64() {

src/tensor/mod.rs

@@ -638,42 +638,44 @@ impl<T: Clone + TensorType> Tensor<T> {
     where
         T: Send + Sync,
     {
-        if indices.is_empty() {
+        // Fast path: empty indices or full tensor slice
+        if indices.is_empty()
+            || indices.iter().map(|x| x.end - x.start).collect::<Vec<_>>() == self.dims
+        {
             return Ok(self.clone());
         }
+
+        // Validate dimensions
         if self.dims.len() < indices.len() {
             return Err(TensorError::DimError(format!(
                 "The dimensionality of the slice {:?} is greater than the tensor's {:?}",
                 indices, self.dims
             )));
-        } else if indices.iter().map(|x| x.end - x.start).collect::<Vec<_>>() == self.dims {
-            // else if slice is the same as dims, return self
-            return Ok(self.clone());
         }
 
-        // if indices weren't specified we fill them in as required
-        let mut full_indices = indices.to_vec();
+        // Pre-allocate the full indices vector with capacity
+        let mut full_indices = Vec::with_capacity(self.dims.len());
+        full_indices.extend_from_slice(indices);
 
-        for i in 0..(self.dims.len() - indices.len()) {
-            full_indices.push(0..self.dims()[indices.len() + i])
-        }
+        // Fill remaining dimensions
+        full_indices.extend((indices.len()..self.dims.len()).map(|i| 0..self.dims[i]));
 
-        let cartesian_coord: Vec<Vec<usize>> = full_indices
+        // Pre-calculate total size and allocate result vector
+        let total_size: usize = full_indices
             .iter()
-            .cloned()
-            .multi_cartesian_product()
-            .collect();
-
-        let res: Vec<T> = cartesian_coord
-            .par_iter()
-            .map(|e| {
-                let index = self.get_index(e);
-                self[index].clone()
-            })
-            .collect();
+            .map(|range| range.end - range.start)
+            .product();
+        let mut res = Vec::with_capacity(total_size);
 
+        // Calculate new dimensions once
         let dims: Vec<usize> = full_indices.iter().map(|e| e.end - e.start).collect();
 
+        // Use iterator directly without collecting into intermediate Vec
+        for coord in full_indices.iter().cloned().multi_cartesian_product() {
+            let index = self.get_index(&coord);
+            res.push(self[index].clone());
+        }
+
         Tensor::new(Some(&res), &dims)
     }