Repository: Lagrange-Labs/deep-prove
Pull request #77: Cleanup tensors.rs
Author: login: sshravan
Created: 2025-03-19T20:07:57Z | Merged: 2025-03-20T13:32:25Z



Diff:
diff --git a/zkml/src/commit/precommit.rs b/zkml/src/commit/precommit.rs
index b1462153..e6ba311c 100644
--- a/zkml/src/commit/precommit.rs
+++ b/zkml/src/commit/precommit.rs
@@ -427,7 +427,7 @@ mod test {
                     rng.gen_range(3..24) as usize,
                     rng.gen_range(3..24) as usize,
                 ])
-                .pad_next_power_of_two_2d()
+                .pad_next_power_of_two()
             })
             .enumerate()
             .collect_vec();
diff --git a/zkml/src/iop/mod.rs b/zkml/src/iop/mod.rs
index cfa58371..78e5878d 100644
--- a/zkml/src/iop/mod.rs
+++ b/zkml/src/iop/mod.rs
@@ -113,8 +113,8 @@ mod test {
         model.describe();
         let trace = model.run(input.clone());
         let output = trace.final_output();
-        let ctx =
-            Context::<F>::generate(&model, Some(input.dims())).expect("unable to generate context");
+        let ctx = Context::<F>::generate(&model, Some(input.get_shape()))
+            .expect("unable to generate context");
         let io = IO::new(input.to_fields(), output.clone().to_fields());
         let mut prover_transcript = default_transcript();
         let prover = Prover::<_, _>::new(&ctx, &mut prover_transcript);
diff --git a/zkml/src/layers/activation.rs b/zkml/src/layers/activation.rs
index e253315f..fd30dfe0 100644
--- a/zkml/src/layers/activation.rs
+++ b/zkml/src/layers/activation.rs
@@ -191,7 +191,7 @@ impl Relu {
 
     pub fn op(&self, input: &Tensor<Element>) -> Tensor<Element> {
         Tensor::new(
-            input.dims(),
+            input.get_shape(),
             input
                 .get_data()
                 .par_iter()
@@ -258,7 +258,7 @@ mod test {
         assert_eq!(input_mle.num_vars(), Relu::num_vars());
         let inputs = Tensor::random(vec![10]);
         let outputs = relu.op(&inputs);
-        assert_eq!(inputs.dims(), outputs.dims());
+        assert_eq!(inputs.get_shape(), outputs.get_shape());
         for (input, output) in inputs.get_data().iter().zip(outputs.get_data().iter()) {
             // here putting input works because every random input is a u8, so it's already within [0;256] so
             // its value "is" the index. Normally if this is not true, we should get the index of the row corresponding to that input
diff --git a/zkml/src/layers/convolution.rs b/zkml/src/layers/convolution.rs
index de7d6f5b..987e940b 100644
--- a/zkml/src/layers/convolution.rs
+++ b/zkml/src/layers/convolution.rs
@@ -71,7 +71,7 @@ pub struct ConvProof<E: ExtensionField> {
 
 impl Convolution {
     pub fn new(filter: Tensor<Element>, bias: Tensor<Element>) -> Self {
-        assert_eq!(filter.kw(), bias.dims()[0]);
+        assert_eq!(filter.kw(), bias.get_shape()[0]);
         Self { filter, bias }
     }
     pub fn add_bias(&self, conv_out: &Tensor<Element>) -> Tensor<Element> {
@@ -192,7 +192,7 @@ impl Convolution {
         // .map(|&dim| ceil_log2(dim)),
         // )
         // .collect::<Vec<usize>>();
-        let mut filter_shape = self.filter.dims();
+        let mut filter_shape = self.filter.get_shape();
         filter_shape.remove(1);
         aux.last_output_shape = filter_shape;
 
diff --git a/zkml/src/layers/dense.rs b/zkml/src/layers/dense.rs
index e0bfdfd9..4aaf83e5 100644
--- a/zkml/src/layers/dense.rs
+++ b/zkml/src/layers/dense.rs
@@ -55,7 +55,7 @@ pub struct DenseProof<E: ExtensionField> {
 
 impl Dense {
     pub fn new(matrix: Tensor<Element>, bias: Tensor<Element>) -> Self {
-        assert_eq!(matrix.nrows_2d(), bias.dims()[0]);
+        assert_eq!(matrix.nrows_2d(), bias.get_shape()[0]);
         Self { matrix, bias }
     }
     pub fn ncols(&self) -> usize {
@@ -66,7 +66,7 @@ impl Dense {
     }
 
     pub fn op(&self, input: &Tensor<Element>) -> Tensor<Element> {
-        if input.dims().len() != 1 {
+        if input.get_shape().len() != 1 {
             let flat_input = input.flatten();
             self.matrix.matvec(&flat_input).add(&self.bias)
         } else {
@@ -75,7 +75,7 @@ impl Dense {
     }
 
     pub fn pad_next_power_of_two(self) -> Self {
-        let matrix = self.matrix.pad_next_power_of_two_2d();
+        let matrix = self.matrix.pad_next_power_of_two();
         let bias = self.bias.pad_1d(matrix.nrows_2d());
         Self::new(matrix, bias)
     }
@@ -376,7 +376,7 @@ mod test {
         fn test_dense_pad_next_power_of_two() {
             // Create a Dense layer with non-power-of-two dimensions
             let matrix =
-                Tensor::<Element>::from_coeffs_2d(vec![vec![1, 2, 3], vec![4, 5, 6], vec![
+                Tensor::<Element>::matix_from_coeffs(vec![vec![1, 2, 3], vec![4, 5, 6], vec![
                     7, 8, 9,
                 ]])
                 .unwrap();
@@ -389,12 +389,12 @@ mod test {
             let padded = dense.pad_next_power_of_two();
 
             // Check padded dimensions are powers of two
-            let padded_dims = padded.matrix.dims();
+            let padded_dims = padded.matrix.get_shape();
             assert_eq!(padded_dims[0], 4); // Next power of 2 after 3
             assert_eq!(padded_dims[1], 4); // Next power of 2 after 3
 
             // Check bias is padded
-            let bias_dims = padded.bias.dims();
+            let bias_dims = padded.bias.get_shape();
             assert_eq!(bias_dims[0], 4); // Next power of 2 after 3
 
             // Check original values are preserved
@@ -419,7 +419,7 @@ mod test {
         #[test]
         fn test_dense_pad_already_power_of_two() {
             // Create a Dense layer with power-of-two dimensions
-            let matrix = Tensor::<Element>::from_coeffs_2d(vec![
+            let matrix = Tensor::<Element>::matix_from_coeffs(vec![
                 vec![1, 2, 3, 4],
                 vec![5, 6, 7, 8],
                 vec![9, 10, 11, 12],
@@ -435,12 +435,12 @@ mod test {
             let padded = dense.clone().pad_next_power_of_two();
 
             // Check dimensions remain the same
-            let padded_dims = padded.matrix.dims();
+            let padded_dims = padded.matrix.get_shape();
             assert_eq!(padded_dims[0], 4);
             assert_eq!(padded_dims[1], 4);
 
             // Check bias dimensions remain the same
-            let bias_dims = padded.bias.dims();
+            let bias_dims = padded.bias.get_shape();
             assert_eq!(bias_dims[0], 4);
 
             // Check values are preserved
@@ -456,11 +456,12 @@ mod test {
         #[test]
         fn test_dense_pad_mixed_dimensions() {
             // Create a Dense layer with one power-of-two dimension and one non-power-of-two
-            let matrix =
-                Tensor::<Element>::from_coeffs_2d(vec![vec![1, 2, 3, 4], vec![5, 6, 7, 8], vec![
-                    9, 10, 11, 12,
-                ]])
-                .unwrap();
+            let matrix = Tensor::<Element>::matix_from_coeffs(vec![
+                vec![1, 2, 3, 4],
+                vec![5, 6, 7, 8],
+                vec![9, 10, 11, 12],
+            ])
+            .unwrap();
 
             let bias = Tensor::<Element>::new(vec![3], vec![20, 21, 22]);
 
@@ -470,12 +471,12 @@ mod test {
             let padded = dense.pad_next_power_of_two();
 
             // Check dimensions are padded correctly
-            let padded_dims = padded.matrix.dims();
+            let padded_dims = padded.matrix.get_shape();
             assert_eq!(padded_dims[0], 4); // Next power of 2 after 3
             assert_eq!(padded_dims[1], 4); // Already a power of 2
 
             // Check bias is padded
-            let bias_dims = padded.bias.dims();
+            let bias_dims = padded.bias.get_shape();
             assert_eq!(bias_dims[0], 4); // Next power of 2 after 3
 
             // Check original values are preserved and padding is zeros
@@ -504,7 +505,7 @@ mod test {
 
             // Create a Dense layer
             let matrix =
-                Tensor::<Element>::from_coeffs_2d(vec![vec![1, 2, 3], vec![4, 5, 6]]).unwrap();
+                Tensor::<Element>::matix_from_coeffs(vec![vec![1, 2, 3], vec![4, 5, 6]]).unwrap();
 
             let bias = Tensor::<Element>::new(vec![2], vec![10, 11]);
 
diff --git a/zkml/src/layers/pooling.rs b/zkml/src/layers/pooling.rs
index 617576e7..719df3f7 100644
--- a/zkml/src/layers/pooling.rs
+++ b/zkml/src/layers/pooling.rs
@@ -142,7 +142,7 @@ impl Pooling {
         E::BaseField: Serialize + DeserializeOwned,
         E: Serialize + DeserializeOwned,
     {
-        assert_eq!(input.dims().len(), 3, "Maxpool needs 3D inputs.");
+        assert_eq!(input.get_shape().len(), 3, "Maxpool needs 3D inputs.");
         // Create the range check proof for the diff
         let prover_info = prover.next_lookup_witness()?;
 
@@ -223,7 +223,7 @@ impl Pooling {
         // Now we must do the samething accumulating evals for the input poly as we fix variables on the input poly.
         // The point length is 2 longer because for now we only support MaxPool2D.
 
-        let padded_input_shape = input.dims();
+        let padded_input_shape = input.get_shape();
         let padded_input_row_length_log = ceil_log2(padded_input_shape[2]);
         // We can batch all of the claims for the input poly with 00, 10, 01, 11 fixed into one with random challenges
         let [r1, r2] = [prover
@@ -446,7 +446,7 @@ impl Maxpool2D {
         let padded_input = input.pad_next_power_of_two();
 
         let padded_output = self.op(&input).pad_next_power_of_two();
-        let padded_input_shape = padded_input.dims();
+        let padded_input_shape = padded_input.get_shape();
 
         let new_fixed = (0..padded_input_shape[2] << 1)
             .into_par_iter()
@@ -573,7 +573,7 @@ mod tests {
 
             let padded_output = output.pad_next_power_of_two();
 
-            let padded_input_shape = padded_input.dims();
+            let padded_input_shape = padded_input.get_shape();
 
             let num_vars = padded_input.get_data().len().ilog2() as usize;
             let output_num_vars = padded_output.get_data().len().ilog2() as usize;
diff --git a/zkml/src/layers/requant.rs b/zkml/src/layers/requant.rs
index bf66c1c3..bc077069 100644
--- a/zkml/src/layers/requant.rs
+++ b/zkml/src/layers/requant.rs
@@ -61,7 +61,7 @@ where
 impl Requant {
     pub fn op(&self, input: &crate::tensor::Tensor<Element>) -> crate::tensor::Tensor<Element> {
         crate::tensor::Tensor::<Element>::new(
-            input.dims(),
+            input.get_shape(),
             input.get_data().iter().map(|e| self.apply(e)).collect_vec(),
         )
     }
diff --git a/zkml/src/model.rs b/zkml/src/model.rs
index 0d167304..668390a5 100644
--- a/zkml/src/model.rs
+++ b/zkml/src/model.rs
@@ -777,7 +777,7 @@ pub(crate) mod test {
         let trace: crate::model::InferenceTrace<'_, _, GoldilocksExt2> =
             model.run::<F>(input.clone());
         let mut tr: BasicTranscript<GoldilocksExt2> = BasicTranscript::new(b"m2vec");
-        let ctx = Context::<GoldilocksExt2>::generate(&model, Some(input.dims()))
+        let ctx = Context::<GoldilocksExt2>::generate(&model, Some(input.get_shape()))
             .expect("Unable to generate context");
         let output = trace.final_output().clone();
 
@@ -828,8 +828,9 @@ pub(crate) mod test {
                             model.run::<F>(input.clone());
                         let mut tr: BasicTranscript<GoldilocksExt2> =
                             BasicTranscript::new(b"m2vec");
-                        let ctx = Context::<GoldilocksExt2>::generate(&model, Some(input.dims()))
-                            .expect("Unable to generate context");
+                        let ctx =
+                            Context::<GoldilocksExt2>::generate(&model, Some(input.get_shape()))
+                                .expect("Unable to generate context");
                         let output = trace.final_output().clone();
                         let prover: Prover<'_, GoldilocksExt2, BasicTranscript<GoldilocksExt2>> =
                             Prover::new(&ctx, &mut tr);
diff --git a/zkml/src/onnx_parse.rs b/zkml/src/onnx_parse.rs
index 5077a7d5..647d78b1 100644
--- a/zkml/src/onnx_parse.rs
+++ b/zkml/src/onnx_parse.rs
@@ -347,9 +347,9 @@ pub fn load_model<Q: Quantizer<Element>>(filepath: &str) -> Result<Model> {
                 )?;
                 let bias =
                     fetch_weight_bias_as_tensor::<Q>("bias", node, &initializers, global_max_abs)?;
-                ensure!(bias.dims().len() == 1, "bias is not a vector");
-                input_shape_og = vec![weight.dims()[0]];
-                let nrows = weight.dims()[0];
+                ensure!(bias.get_shape().len() == 1, "bias is not a vector");
+                input_shape_og = vec![weight.get_shape()[0]];
+                let nrows = weight.get_shape()[0];
                 ensure!(
                     bias.get_data().len() == nrows,
                     "bias length {} does not match matrix width {}",
@@ -394,7 +394,7 @@ pub fn load_model<Q: Quantizer<Element>>(filepath: &str) -> Result<Model> {
                 let bias = bias.pad_1d(nrows);
                 input_shape_padded = vec![nrows];
 
-                debug!("layer idx {} -> final shape {:?}", i, weight.dims());
+                debug!("layer idx {} -> final shape {:?}", i, weight.get_shape());
                 layers.push(Layer::Dense(Dense::new(weight, bias)));
             }
             op if ACTIVATION.contains(&op) => {
@@ -413,13 +413,13 @@ pub fn load_model<Q: Quantizer<Element>>(filepath: &str) -> Result<Model> {
                 let mut bias =
                     fetch_weight_bias_as_tensor::<Q>("bias", node, &initializers, global_max_abs)?;
 
-                input_shape_og = conv2d_shape(&input_shape_og, &weight.dims());
-                let weight_shape = weight.dims();
+                input_shape_og = conv2d_shape(&input_shape_og, &weight.get_shape());
+                let weight_shape = weight.get_shape();
                 // Perform basic sanity checks on the tensor dimensions
                 let shape_test = check_filter(&weight_shape);
                 assert!(shape_test.is_ok(), "Failed: {:?}", shape_test.unwrap_err());
                 assert!(
-                    weight_shape[0] == bias.dims()[0],
+                    weight_shape[0] == bias.get_shape()[0],
                     "Bias length doesn't match filter shape"
                 );
 
@@ -432,7 +432,7 @@ pub fn load_model<Q: Quantizer<Element>>(filepath: &str) -> Result<Model> {
                 assert!(input_shape_padded.len() == 3);
 
                 // Since we are doing an FFT based conv, we need to pad the last two dimensions of the filter to match the input.
-                let weight_shape = weight.dims();
+                let weight_shape = weight.get_shape();
                 let (filter_height, filter_weight) = (weight_shape[2], weight_shape[3]);
                 let (input_height, input_weight) = (input_shape_padded[1], input_shape_padded[2]);
 
@@ -446,9 +446,9 @@ pub fn load_model<Q: Quantizer<Element>>(filepath: &str) -> Result<Model> {
                 // Filter need to know the shape of the input
                 // weight.update_input_shape(&input_shape_padded);
 
-                let dims = weight.dims();
+                let dims = weight.get_shape();
                 let weight = crate::tensor::Tensor::new_conv(
-                    weight.dims(),
+                    weight.get_shape(),
                     input_shape_padded.clone(),
                     weight.get_data().to_vec(),
                 );
@@ -853,7 +853,7 @@ mod tests {
         // println!("Result: {:?}", trace.final_output());
 
         let mut tr: BasicTranscript<GoldilocksExt2> = BasicTranscript::new(b"m2vec");
-        let ctx = Context::<GoldilocksExt2>::generate(&model, Some(input.dims()))
+        let ctx = Context::<GoldilocksExt2>::generate(&model, Some(input.get_shape()))
             .expect("Unable to generate context");
         let output = trace.final_output().clone();
 
diff --git a/zkml/src/quantization.rs b/zkml/src/quantization.rs
index 5607aacf..c915398a 100644
--- a/zkml/src/quantization.rs
+++ b/zkml/src/quantization.rs
@@ -122,7 +122,7 @@ where
 {
     fn to_fields(self) -> Tensor<F> {
         Tensor::new(
-            self.dims(),
+            self.get_shape(),
             self.get_data()
                 .into_iter()
                 .map(|i| i.to_field())
diff --git a/zkml/src/tensor.rs b/zkml/src/tensor.rs
index 03297993..affafb5d 100644
--- a/zkml/src/tensor.rs
+++ b/zkml/src/tensor.rs
@@ -1,6 +1,4 @@
 use anyhow::bail;
-use ark_std::rand::{self, SeedableRng, rngs::StdRng};
-use ff::Field;
 use ff_ext::ExtensionField;
 use goldilocks::GoldilocksExt2;
 use itertools::Itertools;
@@ -10,6 +8,7 @@ use rayon::{
         IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator,
         IntoParallelRefMutIterator, ParallelIterator,
     },
+    prelude::ParallelSlice,
     slice::ParallelSliceMut,
 };
 use std::{
@@ -49,6 +48,12 @@ pub fn check_tensor_consistency(real_tensor: Tensor<Element>, padded_tensor: Ten
     }
 }
 
+/// Returns an n-th root of unity by starting with a 32nd root of unity and squaring it (32-n) times.
+/// Each squaring operation halves the order of the root of unity:
+/// - For n=16: squares it 16 times (32-16) to get a 16th root of unity
+/// - For n=8:  squares it 24 times (32-8) to get an 8th root of unity
+/// - For n=4:  squares it 28 times (32-4) to get a 4th root of unity
+/// The initial ROOT_OF_UNITY constant is verified to be a 32nd root of unity in the field implementation.
 pub fn get_root_of_unity<E: ExtensionField>(n: usize) -> E {
     let mut rou = E::ROOT_OF_UNITY;
 
@@ -58,50 +63,47 @@ pub fn get_root_of_unity<E: ExtensionField>(n: usize) -> E {
 
     return rou;
 }
-
 // Properly pad a filter
-pub fn index_w<E: ExtensionField>(w: Vec<Element>, vec: &mut Vec<E>, n_real: usize, n: usize) {
-    // let mut vec = vec![E::ZERO;n*n];
-    for i in 0..n_real {
-        for j in 0..n_real {
-            if w[i * n_real + j] < 0 {
-                vec[i * n + j] = -E::from((0 - w[i * n_real + j]) as u64);
-            } else {
-                vec[i * n + j] = E::from((w[i * n_real + j]) as u64);
-            }
+pub fn index_w<E: ExtensionField>(
+    w: &[Element],
+    n_real: usize,
+    n: usize,
+    output_len: usize,
+) -> impl ParallelIterator<Item = E> + use<'_, E> {
+    (0..output_len).into_par_iter().map(move |idx| {
+        let i = idx / n;
+        let j = idx % n;
+        if i < n_real && j < n_real {
+            w[i * n_real + j].to_field()
+        } else {
+            E::ZERO
         }
-    }
+    })
 }
-
 // let u = [u[1],...,u[n*n]]
 // output vec = [u[n*n-1],u[n*n-2],...,u[n*n-n],....,u[0]]
 // Note that y_eval =  f_vec(r) = f_u(1-r)
 pub fn index_u<E: ExtensionField>(u: Vec<E>, n: usize) -> Vec<E> {
     let mut vec = vec![E::ZERO; u.len() / 2];
-    for i in 0..n {
-        for j in 0..n {
-            vec[i * n + j] = u[n * n - 1 - i * n - j];
-        }
-    }
+    let len = n * n;
+    vec.par_iter_mut().enumerate().for_each(|(i, val)| {
+        *val = u[len - 1 - i];
+    });
     return vec;
 }
-
 // let x: [x[0][0],...,x[0][n],x[1][0],...,x[n][n]]
 // output vec = [x[n][n], x[n][n-1],...,x[n][0],x[n-1]x[n],...,x[0][0]]
 // Note that y_eval = f_vec(r) = f_x(1-r)
 pub fn index_x<E: ExtensionField>(x: Vec<Element>, vec: &mut Vec<E>, n: usize) {
-    for i in 0..n {
-        for j in 0..n {
-            let val = x[n * (n - 1 - i) + n - 1 - j];
-            vec[i * n + j] = val.to_field();
-        }
-    }
+    let len = n * n;
+    vec.par_iter_mut().enumerate().for_each(|(i, val)| {
+        *val = x[len - 1 - i].to_field();
+    });
 }
-
 // FFT implementation,
 // flag: false -> FFT
 // flag: true -> iFFT
-pub fn fft<E: ExtensionField>(v: &mut Vec<E>, flag: bool) {
+pub fn fft<E: ExtensionField + Send + Sync>(v: &mut Vec<E>, flag: bool) {
     let n = v.len();
     let logn = ark_std::log2(n) as u32;
     let mut rev: Vec<usize> = vec![0; n];
@@ -124,25 +126,30 @@ pub fn fft<E: ExtensionField>(v: &mut Vec<E>, flag: bool) {
     for i in 2..n {
         w[i] = w[i - 1] * w[1];
     }
-    for i in 0..n {
-        if rev[i] < (i) {
-            let temp = v[i];
-            v[i] = v[rev[i]];
-            v[rev[i]] = temp;
-        }
+
+    // Collect indices that need to be swapped
+    let swaps: Vec<(usize, usize)> = (0..n)
+        .into_par_iter()
+        .filter_map(|i| if rev[i] < i { Some((i, rev[i])) } else { None })
+        .collect();
+
+    // Perform swaps sequentially
+    for (i, j) in swaps {
+        v.swap(i, j);
     }
-    let mut u: E;
-    let mut l: E;
+
     let mut i: usize = 2;
     while i <= n {
-        for j in (0..n).step_by(i) {
-            for k in 0..(i >> 1) {
-                u = v[j + k];
-                l = v[j + k + (i >> 1)] * w[n / i * k];
-                v[j + k] = u + l;
-                v[j + k + (i >> 1)] = u - l;
+        // Parallelize the FFT butterfly operations
+        v.par_chunks_mut(i).for_each(|chunk| {
+            let half_i = i >> 1;
+            for k in 0..half_i {
+                let u = chunk[k];
+                let l = chunk[k + half_i] * w[n / i * k];
+                chunk[k] = u + l;
+                chunk[k + half_i] = u - l;
             }
-        }
+        });
         i <<= 1;
     }
 
@@ -152,19 +159,12 @@ pub fn fft<E: ExtensionField>(v: &mut Vec<E>, flag: bool) {
         if ilen * E::from(n as u64) != E::ONE {
             println!("Error in inv\n");
         }
-        for i in 0..n {
-            v[i] = v[i] * ilen;
-        }
+        v.par_iter_mut().for_each(|val| {
+            *val = *val * ilen;
+        });
     }
 }
 
-#[derive(Debug, Clone)]
-pub struct Tensor<T> {
-    pub data: Vec<T>,
-    pub shape: Vec<usize>,
-    input_shape: Vec<usize>,
-}
-
 pub struct ConvData<E>
 where
     E: Clone + ExtensionField,
@@ -212,6 +212,13 @@ where
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct Tensor<T> {
+    pub data: Vec<T>,
+    pub shape: Vec<usize>,
+    input_shape: Vec<usize>,
+}
+
 impl Tensor<Element> {
     // Create specifically a new convolution. The input shape is needed to compute the
     // output and properly arrange the weights
@@ -222,52 +229,29 @@ impl Tensor<Element> {
         );
         assert!(shape.len() == 4, "Shape does not match data length.");
 
+        let i0 = input_shape[0];
+        let s2 = shape[2];
+        let rdata = &data[..];
         let n_w = (input_shape[1] - shape[2] + 1).next_power_of_two();
-        let mut w_fft =
-            vec![
-                vec![vec![GoldilocksExt2::ZERO; n_w * n_w]; input_shape[0].next_power_of_two()];
-                shape[0].next_power_of_two()
-            ];
-        for i in 0..shape[0] {
-            for j in 0..input_shape[0] {
-                index_w(
-                    data[(i * input_shape[0] * shape[2] * shape[2] + j * shape[2] * shape[2])
-                        ..(i * input_shape[0] * shape[2] * shape[2]
-                            + (j + 1) * shape[2] * shape[2])]
-                        .to_vec(),
-                    &mut w_fft[i][j],
-                    shape[2],
-                    n_w,
-                );
-                w_fft[i][j].resize(2 * n_w * n_w, GoldilocksExt2::ZERO);
-                fft(&mut w_fft[i][j], false);
-            }
-        }
-        let mut w: Vec<Element> = vec![0; w_fft.len() * w_fft[0].len() * w_fft[0][0].len()];
-        let mut ctr = 0;
-        for i in 0..w_fft.len() {
-            for j in 0..w_fft[0].len() {
-                for k in 0..w_fft[0][0].len() {
-                    if GoldilocksExt2::to_canonical_u64_vec(&w_fft[i][j][k])[0] as u64
-                        > (1 << 60 as u64)
-                    {
-                        w[ctr] = -(GoldilocksExt2::to_canonical_u64_vec(&(-w_fft[i][j][k]))[0]
-                            as Element);
-                    } else {
-                        w[ctr] =
-                            GoldilocksExt2::to_canonical_u64_vec(&(w_fft[i][j][k]))[0] as Element;
-                    }
-                    ctr = ctr + 1;
-                }
-            }
-        }
+        let w_fft = (0..shape[0])
+            .into_par_iter()
+            .flat_map(move |i| {
+                (0..i0).into_par_iter().flat_map(move |j| {
+                    let range =
+                        (i * i0 * s2 * s2 + j * s2 * s2)..(i * i0 * s2 * s2 + (j + 1) * s2 * s2);
+                    let mut w = index_w(&rdata[range], s2, n_w, 2 * n_w * n_w)
+                        .collect::<Vec<GoldilocksExt2>>();
+                    fft(&mut w, false);
+                    w.into_par_iter().map(|e| e.into_element())
+                })
+            })
+            .collect::<Vec<_>>();
         Self {
-            data: w,
+            data: w_fft,
             shape: vec![shape[0], shape[1], n_w, n_w],
             input_shape,
         }
     }
-
     pub fn get_real_weights<F: ExtensionField>(&self) -> Vec<Vec<Vec<Element>>> {
         let mut w_fft =
             vec![
@@ -278,11 +262,7 @@ impl Tensor<Element> {
         for i in 0..w_fft.len() {
             for j in 0..w_fft[i].len() {
                 for k in 0..w_fft[i][j].len() {
-                    if self.data[ctr] < 0 {
-                        w_fft[i][j][k] = -F::from((-self.data[ctr]) as u64);
-                    } else {
-                        w_fft[i][j][k] = F::from((self.data[ctr]) as u64);
-                    }
+                    w_fft[i][j][k] = self.data[ctr].to_field();
                     ctr += 1;
                 }
             }
@@ -297,20 +277,12 @@ impl Tensor<Element> {
         for i in 0..self.kw() {
             for j in 0..self.kx() {
                 for k in 0..(self.nw() * self.nw()) {
-                    if F::to_canonical_u64_vec(&w_fft[i][j][k])[0] as u64 > (1 << 60 as u64) {
-                        real_weights[i][j][k] =
-                            //-(F::to_canonical_u64_vec(&(-w_fft[i][j][k]))[0] as Element);
-                            w_fft[i][j][k].into_element();
-                    } else {
-                        real_weights[i][j][k] = w_fft[i][j][k].into_element();
-                        // F::to_canonical_u64_vec(&(w_fft[i][j][k]))[0] as Element;
-                    }
+                    real_weights[i][j][k] = w_fft[i][j][k].into_element();
                 }
             }
         }
         real_weights
     }
-
     // Convolution algorithm using FFTs.
     // When invoking this algorithm the prover generates all withness/intermidiate evaluations
     // needed to generate a convolution proof
@@ -320,19 +292,15 @@ impl Tensor<Element> {
     ) -> (Tensor<Element>, ConvData<F>) {
         let n_x = x.shape[1].next_power_of_two();
         let mut real_input = vec![F::ZERO; x.data.len()];
-        for i in 0..real_input.len() {
-            if x.data[i] < 0 {
-                real_input[i] = -F::from((-x.data[i]) as u64);
-            } else {
-                real_input[i] = F::from(x.data[i] as u64);
-            }
-        }
+        real_input.par_iter_mut().enumerate().for_each(|(i, val)| {
+            *val = x.data[i].to_field();
+        });
 
         let mut x_vec = vec![vec![F::ZERO; n_x * n_x]; x.shape[0].next_power_of_two()];
         let mut w_fft = vec![F::ZERO; self.data.len()];
-        for i in 0..w_fft.len() {
-            w_fft[i] = self.data[i].to_field();
-        }
+        w_fft.par_iter_mut().enumerate().for_each(|(i, val)| {
+            *val = self.data[i].to_field();
+        });
 
         for i in 0..x_vec.len() {
             index_x(
@@ -374,13 +342,6 @@ impl Tensor<Element> {
             for j in 0..out[i].len() {
                 let val = out[i][j].into_element();
                 out_element[i * out[i].len() + j] = val;
-                // if F::to_canonical_u64_vec(&out[i][j])[0] as u64 > (1 << 60 as u64) {
-                //    out_element[i * out[i].len() + j] =
-                //        -(F::to_canonical_u64_vec(&(-out[i][j]))[0] as Element);
-                //} else {
-                //    out_element[i * out[i].len() + j] =
-                //        F::to_canonical_u64_vec(&(out[i][j]))[0] as Element;
-                //}
             }
         }
 
@@ -389,7 +350,6 @@ impl Tensor<Element> {
             ConvData::new(real_input, input, input_fft, prod, output),
         );
     }
-
     pub fn kx(&self) -> usize {
         self.input_shape[0]
     }
@@ -403,6 +363,38 @@ impl Tensor<Element> {
     pub fn filter_size(&self) -> usize {
         self.shape[2] * self.shape[2]
     }
+    /// Returns the evaluation point, in order for (row,col) addressing
+    pub fn evals_2d<F: ExtensionField>(&self) -> Vec<F> {
+        assert!(self.is_matrix(), "Tensor is not a matrix");
+        self.evals_flat()
+    }
+    pub fn evals_flat<F: ExtensionField>(&self) -> Vec<F> {
+        self.data.par_iter().map(|e| e.to_field()).collect()
+    }
+    pub fn get_conv_weights<F: ExtensionField>(&self) -> Vec<F> {
+        let mut data = vec![F::ZERO; self.data.len()];
+        for i in 0..data.len() {
+            data[i] = self.data[i].to_field();
+        }
+        data
+    }
+    /// Returns a MLE of the matrix that can be evaluated.
+    pub fn to_mle_2d<F: ExtensionField>(&self) -> DenseMultilinearExtension<F> {
+        assert!(self.is_matrix(), "Tensor is not a matrix");
+        assert!(
+            self.nrows_2d().is_power_of_two(),
+            "number of rows {} is not a power of two",
+            self.nrows_2d()
+        );
+        assert!(
+            self.ncols_2d().is_power_of_two(),
+            "number of columns {} is not a power of two",
+            self.ncols_2d()
+        );
+        // N variable to address 2^N rows and M variables to address 2^M columns
+        let num_vars = self.nrows_2d().ilog2() + self.ncols_2d().ilog2();
+        DenseMultilinearExtension::from_evaluations_ext_vec(num_vars as usize, self.evals_2d())
+    }
 }
 
 impl<T> Tensor<T> {
@@ -418,35 +410,22 @@ impl<T> Tensor<T> {
             input_shape: vec![0],
         }
     }
-
-    /// Get the dimensions of the tensor
-    pub fn dims(&self) -> Vec<usize> {
-        assert!(self.shape.len() > 0, "Empty tensor");
-        self.shape.clone()
-    }
-
-    /// Get the dimensions of the tensor
-    pub fn get_input_shape(&self) -> Vec<usize> {
-        assert!(self.shape.len() > 0, "Empty tensor");
-        self.input_shape.clone()
-    }
-
     /// Is vector
     pub fn is_vector(&self) -> bool {
-        self.dims().len() == 1
-    }
-    pub fn is_convolution(&self) -> bool {
-        self.dims().len() == 4
+        self.get_shape().len() == 1
     }
     /// Is matrix
     pub fn is_matrix(&self) -> bool {
-        self.dims().len() == 2
+        self.get_shape().len() == 2
+    }
+    ///
+    pub fn is_convolution(&self) -> bool {
+        self.get_shape().len() == 4
     }
-
     /// Get the number of rows from the matrix
     pub fn nrows_2d(&self) -> usize {
         let mut cols = 0;
-        let dims = self.dims();
+        let dims = self.get_shape();
         if self.is_matrix() {
             cols = dims[0];
         } else if self.is_convolution() {
@@ -455,11 +434,10 @@ impl<T> Tensor<T> {
         assert!(cols != 0, "Tensor is not a matrix or convolution");
         cols
     }
-
     /// Get the number of cols from the matrix
     pub fn ncols_2d(&self) -> usize {
         let mut cols = 0;
-        let dims = self.dims();
+        let dims = self.get_shape();
         if self.is_matrix() {
             cols = dims[1];
         } else if self.is_convolution() {
@@ -471,7 +449,6 @@ impl<T> Tensor<T> {
 
         return cols;
     }
-
     /// Returns the number of boolean variables needed to address any row, and any columns
     pub fn num_vars_2d(&self) -> (usize, usize) {
         assert!(self.is_matrix(), "Tensor is not a matrix");
@@ -480,12 +457,22 @@ impl<T> Tensor<T> {
             self.ncols_2d().ilog2() as usize,
         )
     }
-
+    /// Get the dimensions of the tensor
+    pub fn get_shape(&self) -> Vec<usize> {
+        assert!(self.shape.len() > 0, "Empty tensor");
+        self.shape.clone()
+    }
+    /// Get the input shape of the tensor
+    /// TODO: Remove it
+    pub fn get_input_shape(&self) -> Vec<usize> {
+        assert!(self.shape.len() > 0, "Empty tensor");
+        self.input_shape.clone()
+    }
     ///
     pub fn get_data(&self) -> &[T] {
         &self.data
     }
-
+    /// TODO: Delete
     pub fn update_input_shape(&mut self, input: &[usize]) {
         // TODO: Need to assert the input_shape with self.shape
         self.input_shape = input.to_vec();
@@ -494,10 +481,7 @@ impl<T> Tensor<T> {
 
 impl<T> Tensor<T>
 where
-    T: Copy + Clone + Send + Sync,
-    T: std::iter::Sum,
-    T: std::ops::Add<Output = T> + std::ops::Sub<Output = T> + std::ops::Mul<Output = T>,
-    T: std::default::Default,
+    T: Clone,
 {
     ///
     pub fn flatten(&self) -> Self {
@@ -509,68 +493,7 @@ where
             input_shape: vec![0],
         }
     }
-
-    pub fn get_shape(&self) -> Vec<usize> {
-        return self.shape.clone();
-    }
-    /// Element-wise addition
-    pub fn add(&self, other: &Tensor<T>) -> Tensor<T> {
-        assert!(self.shape == other.shape, "Shape mismatch for addition.");
-        Tensor {
-            shape: self.shape.clone(),
-            input_shape: vec![0],
-            data: self
-                .data
-                .iter()
-                .zip(other.data.iter())
-                .map(|(a, b)| *a + *b)
-                .collect(),
-        }
-    }
-
-    /// Element-wise subtraction
-    pub fn sub(&self, other: &Tensor<T>) -> Tensor<T> {
-        assert!(self.shape == other.shape, "Shape mismatch for subtraction.");
-        Tensor {
-            shape: self.shape.clone(),
-            input_shape: vec![0],
-            data: self
-                .data
-                .iter()
-                .zip(other.data.iter())
-                .map(|(a, b)| *a - *b)
-                .collect(),
-        }
-    }
-
-    /// Element-wise multiplication
-    pub fn mul(&self, other: &Tensor<T>) -> Tensor<T> {
-        assert!(
-            self.shape == other.shape,
-            "Shape mismatch for multiplication."
-        );
-        Tensor {
-            shape: self.shape.clone(),
-            input_shape: vec![0],
-            data: self
-                .data
-                .iter()
-                .zip(other.data.iter())
-                .map(|(a, b)| *a * *b)
-                .collect(),
-        }
-    }
-
-    /// Scalar multiplication
-    pub fn scalar_mul(&self, scalar: &T) -> Tensor<T> {
-        Tensor {
-            shape: self.shape.clone(),
-            input_shape: vec![0],
-            data: self.data.iter().map(|x| *x * *scalar).collect(),
-        }
-    }
-
-    pub fn from_coeffs_2d(data: Vec<Vec<T>>) -> anyhow::Result<Self> {
+    pub fn matix_from_coeffs(data: Vec<Vec<T>>) -> anyhow::Result<Self> {
         let n_rows = data.len();
         let n_cols = data.first().expect("at least one row in a matrix").len();
         let data = data.into_iter().flatten().collect::<Vec<_>>();
@@ -586,7 +509,6 @@ where
             input_shape: vec![0],
         })
     }
-
     /// Returns the boolean iterator indicating the given row in the right endianness to be
     /// evaluated by an MLE
     pub fn row_to_boolean_2d<F: ExtensionField>(&self, row: usize) -> impl Iterator<Item = F> {
@@ -594,7 +516,6 @@ where
         let (nvars_rows, _) = self.num_vars_2d();
         to_bit_sequence_le(row, nvars_rows).map(|b| F::from(b as u64))
     }
-
     /// Returns the boolean iterator indicating the given row in the right endianness to be
     /// evaluated by an MLE
     pub fn col_to_boolean_2d<F: ExtensionField>(&self, col: usize) -> impl Iterator<Item = F> {
@@ -602,7 +523,6 @@ where
         let (_, nvars_col) = self.num_vars_2d();
         to_bit_sequence_le(col, nvars_col).map(|b| F::from(b as u64))
     }
-
     /// From a given row and a given column, return the vector of field elements in the right
     /// format to evaluate the MLE.
     /// little endian so we need to read cols before rows
@@ -612,7 +532,15 @@ where
             .chain(self.row_to_boolean_2d(row))
             .collect_vec()
     }
+}
 
+impl<T> Tensor<T>
+where
+    T: Copy + Clone + Send + Sync,
+    T: std::iter::Sum,
+    T: std::ops::Add<Output = T> + std::ops::Sub<Output = T> + std::ops::Mul<Output = T>,
+    T: std::default::Default,
+{
     /// Create a tensor filled with zeros
     pub fn zeros(shape: Vec<usize>) -> Self {
         let size = shape.iter().product();
@@ -623,7 +551,59 @@ where
             input_shape: vec![0],
         }
     }
+    /// Element-wise addition
+    pub fn add(&self, other: &Tensor<T>) -> Tensor<T> {
+        assert!(self.shape == other.shape, "Shape mismatch for addition.");
+        let mut data = vec![Default::default(); self.data.len()];
+        data.par_iter_mut().enumerate().for_each(|(i, val)| {
+            *val = self.data[i] + other.data[i];
+        });
+
+        Tensor {
+            shape: self.shape.clone(),
+            input_shape: vec![0],
+            data,
+        }
+    }
+    /// Element-wise subtraction
+    pub fn sub(&self, other: &Tensor<T>) -> Tensor<T> {
+        assert!(self.shape == other.shape, "Shape mismatch for subtraction.");
+        let mut data = vec![Default::default(); self.data.len()];
+        data.par_iter_mut().enumerate().for_each(|(i, val)| {
+            *val = self.data[i] - other.data[i];
+        });
+
+        Tensor {
+            shape: self.shape.clone(),
+            input_shape: vec![0],
+            data,
+        }
+    }
+    /// Element-wise multiplication
+    pub fn mul(&self, other: &Tensor<T>) -> Tensor<T> {
+        assert!(
+            self.shape == other.shape,
+            "Shape mismatch for multiplication."
+        );
+        let mut data = vec![Default::default(); self.data.len()];
+        data.par_iter_mut().enumerate().for_each(|(i, val)| {
+            *val = self.data[i] * other.data[i];
+        });
 
+        Tensor {
+            shape: self.shape.clone(),
+            input_shape: vec![0],
+            data,
+        }
+    }
+    /// Scalar multiplication
+    pub fn scalar_mul(&self, scalar: &T) -> Tensor<T> {
+        Tensor {
+            shape: self.shape.clone(),
+            input_shape: vec![0],
+            data: self.data.par_iter().map(|x| *x * *scalar).collect(),
+        }
+    }
     pub fn pad_1d(mut self, new_len: usize) -> Self {
         assert!(
             self.shape.len() == 1,
@@ -633,8 +613,7 @@ where
         self.shape[0] = new_len;
         self
     }
-
-    pub fn pad_next_power_of_two_2d(mut self) -> Self {
+    fn pad_next_power_of_two_2d(mut self) -> Self {
         assert!(self.is_matrix(), "Tensor is not a matrix");
         // assume the matrix is already well formed and there is always n_rows and n_cols
         // this is because we control the creation of the matrix in the first place
@@ -678,10 +657,9 @@ where
 
         self
     }
-
     /// Recursively pads the tensor so its ready to be viewed as an MLE
     pub fn pad_next_power_of_two(&self) -> Self {
-        let shape = self.dims();
+        let shape = self.get_shape();
 
         let padded_data = Self::recursive_pad(self.get_data(), &shape);
 
@@ -692,7 +670,6 @@ where
 
         Tensor::<T>::new(padded_shape, padded_data)
     }
-
     fn recursive_pad(data: &[T], remaining_dims: &[usize]) -> Vec<T> {
         match remaining_dims.len() {
             // If the remaining dims show we are a vector simply pad
@@ -712,7 +689,7 @@ where
             _ => {
                 let chunk_size = remaining_dims[1..].iter().product::<usize>();
                 let mut unpadded_data = data
-                    .chunks(chunk_size)
+                    .par_chunks(chunk_size)
                     .map(|data_chunk| Self::recursive_pad(data_chunk, &remaining_dims[1..]))
                     .collect::<Vec<Vec<T>>>();
                 let elem_size = unpadded_data[0].len();
@@ -724,51 +701,6 @@ where
             }
         }
     }
-
-    pub fn pad_last_two_dimensions(&self, target: Vec<usize>) -> Self {
-        assert!(self.shape.len() > 2, "Tensor must have 2 dimensions.");
-        assert!(target.len() == 2, "Tensor must have at least 2 dimensions.");
-
-        let (target_x, target_y) = (target[0], target[1]);
-        let current_x = self.shape[self.shape.len() - 2];
-        let current_y = self.shape[self.shape.len() - 1];
-
-        let pad_x = target_x - current_x;
-        let pad_y = target_y - current_y;
-
-        if pad_x == 0 && pad_y == 0 {
-            return self.clone();
-        }
-
-        let mut new_shape = self.shape.clone();
-        new_shape[self.shape.len() - 2] = target_x;
-        new_shape[self.shape.len() - 1] = target_y;
-        let new_size: usize = new_shape.iter().product();
-
-        let mut new_data = vec![T::default(); new_size];
-
-        let mut old_index = 0;
-        let mut new_index = 0;
-
-        // Iterate over all dimensions except the last two
-        let outer_dims: usize = self.shape[..self.shape.len() - 2].iter().product();
-        for _ in 0..outer_dims {
-            // Copy the original rows
-            for row in 0..current_x {
-                for col in 0..current_y {
-                    new_data[new_index + row * target_y + col] =
-                        self.data[old_index + row * current_y + col].clone();
-                }
-            }
-            old_index += current_x * current_y;
-            new_index += target_x * target_y;
-        }
-
-        let mut result = Tensor::new(new_shape, new_data);
-        result.update_input_shape(&self.input_shape);
-        result
-    }
-
     pub fn pad_to_shape(&mut self, target_shape: Vec<usize>) {
         if target_shape.len() != self.shape.len() {
             panic!("Target shape must have the same number of dimensions as the tensor.");
@@ -820,7 +752,6 @@ where
         self.data = new_data;
         self.shape = target_shape;
     }
-
     /// Perform matrix-matrix multiplication
     pub fn matmul(&self, other: &Tensor<T>) -> Tensor<T> {
         assert!(
@@ -854,7 +785,6 @@ where
 
         result
     }
-
     /// Perform matrix-vector multiplication
     /// TODO: actually getting the result should be done via proper tensor-like libraries
     pub fn matvec(&self, vector: &Tensor<T>) -> Tensor<T> {
@@ -877,38 +807,37 @@ where
 
         result
     }
-
     pub fn conv_prod(&self, x: &Vec<Vec<T>>, w: &Vec<Vec<T>>, ii: usize, jj: usize) -> T {
-        let mut sum = Default::default();
-        for i in 0..w.len() {
-            for j in 0..w[i].len() {
-                sum = sum + w[i][j] * x[i + ii][j + jj];
-            }
-        }
-        return sum;
+        w.par_iter()
+            .enumerate()
+            .map(|(i, w_row)| {
+                w_row
+                    .par_iter()
+                    .enumerate()
+                    .map(|(j, &w_val)| w_val * x[i + ii][j + jj])
+                    .sum()
+            })
+            .sum()
     }
-
     pub fn single_naive_conv(&self, w: Vec<Vec<T>>, x: Vec<Vec<T>>) -> Vec<Vec<T>> {
         let mut out: Vec<Vec<T>> =
             vec![vec![Default::default(); x[0].len() - w[0].len() + 1]; x.len() - w.len() + 1];
-        for i in 0..out.len() {
-            for j in 0..out[i].len() {
-                out[i][j] = self.conv_prod(&x, &w, i, j);
-            }
-        }
-        return out;
+        out.par_iter_mut().enumerate().for_each(|(i, out_row)| {
+            out_row.par_iter_mut().enumerate().for_each(|(j, out_val)| {
+                *out_val = self.conv_prod(&x, &w, i, j);
+            });
+        });
+        out
     }
-
     pub fn add_matrix(&self, m1: &mut Vec<Vec<T>>, m2: Vec<Vec<T>>) -> Vec<Vec<T>> {
         let mut m = vec![vec![Default::default(); m1[0].len()]; m1.len()];
-        for i in 0..m.len() {
-            for j in 0..m[i].len() {
-                m[i][j] = m1[i][j] + m2[i][j];
-            }
-        }
+        m.par_iter_mut().enumerate().for_each(|(i, row)| {
+            row.par_iter_mut().enumerate().for_each(|(j, val)| {
+                *val = m1[i][j] + m2[i][j];
+            });
+        });
         return m;
     }
-
     // Implementation of the stadard convolution algorithm.
     // This is needed mostly for debugging purposes
     pub fn cnn_naive_convolution(&self, xt: &Tensor<T>) -> Tensor<T> {
@@ -960,21 +889,24 @@ where
                 .collect(),
         );
     }
-
     /// Transpose the matrix (2D tensor)
     pub fn transpose(&self) -> Tensor<T> {
         assert!(self.is_matrix(), "Tensor is not a matrix.");
         let (m, n) = (self.shape[0], self.shape[1]);
 
         let mut result = Tensor::zeros(vec![n, m]);
-        for i in 0..m {
-            for j in 0..n {
-                result.data[j * m + i] = self.data[i * n + j];
-            }
-        }
         result
-    }
+            .data
+            .par_iter_mut()
+            .enumerate()
+            .for_each(|(idx, val)| {
+                let i = idx % m; // Row in the result matrix
+                let j = idx / m; // Column in the result matrix
+                *val = self.data[i * n + j];
+            });
 
+        result
+    }
     /// Concatenate a matrix (2D tensor) with a vector (1D tensor) as columns
     pub fn concat_matvec_col(&self, vector: &Tensor<T>) -> Tensor<T> {
         assert!(self.is_matrix(), "First tensor is not a matrix.");
@@ -1025,57 +957,20 @@ where
             new_cols
         );
 
-        let mut result = Tensor::<T>::zeros(new_shape);
-
-        // Create a new matrix with expanded dimensions
-        for i in 0..old_rows {
-            for j in 0..old_cols {
-                result.data[i * new_cols + j] = self.data[i * old_cols + j];
-            }
-        }
-        *self = result;
-    }
-}
-
-impl Tensor<Element> {
-    /// Returns the evaluation point, in order for (row,col) addressing
-    pub fn evals_2d<F: ExtensionField>(&self) -> Vec<F> {
-        assert!(self.is_matrix(), "Tensor is not a matrix");
-        self.evals_flat()
-    }
-
-    pub fn evals_flat<F: ExtensionField>(&self) -> Vec<F> {
-        self.data.par_iter().map(|e| e.to_field()).collect()
-    }
-
-    pub fn get_conv_weights<F: ExtensionField>(&self) -> Vec<F> {
-        let mut data = vec![F::ZERO; self.data.len()];
-        for i in 0..data.len() {
-            if self.data[i] < 0 {
-                data[i] = -F::from((-self.data[i]) as u64);
-            } else {
-                data[i] = F::from(self.data[i] as u64);
-            }
-        }
-        data
-    }
+        let new_data: Vec<T> = (0..new_rows * new_cols)
+            .into_par_iter()
+            .map(|idx| {
+                let i = idx / new_cols;
+                let j = idx % new_cols;
+                if i < old_rows && j < old_cols {
+                    self.data[i * old_cols + j].clone()
+                } else {
+                    T::default() // Zero or default for padding
+                }
+            })
+            .collect();
 
-    /// Returns a MLE of the matrix that can be evaluated.
-    pub fn to_mle_2d<F: ExtensionField>(&self) -> DenseMultilinearExtension<F> {
-        assert!(self.is_matrix(), "Tensor is not a matrix");
-        assert!(
-            self.nrows_2d().is_power_of_two(),
-            "number of rows {} is not a power of two",
-            self.nrows_2d()
-        );
-        assert!(
-            self.ncols_2d().is_power_of_two(),
-            "number of columns {} is not a power of two",
-            self.ncols_2d()
-        );
-        // N variable to address 2^N rows and M variables to address 2^M columns
-        let num_vars = self.nrows_2d().ilog2() + self.ncols_2d().ilog2();
-        DenseMultilinearExtension::from_evaluations_ext_vec(num_vars as usize, self.evals_2d())
+        *self = Tensor::new(new_shape, new_data);
     }
 }
 
@@ -1085,7 +980,7 @@ where
     T: std::default::Default,
 {
     pub fn maxpool2d(&self, kernel_size: usize, stride: usize) -> Tensor<T> {
-        let dims = self.dims().len();
+        let dims = self.get_shape().len();
         assert!(dims >= 2, "Input tensor must have at least 2 dimensions.");
 
         let (h, w) = (self.shape[dims - 2], self.shape[dims - 1]);
@@ -1139,14 +1034,14 @@ where
             input_shape: vec![0],
         }
     }
-
+    // Replaces every value of a tensor with the maxpool of its kernel
     pub fn padded_maxpool2d(&self) -> (Tensor<T>, Tensor<T>) {
         let kernel_size = MAXPOOL2D_KERNEL_SIZE;
         let stride = MAXPOOL2D_KERNEL_SIZE;
 
         let maxpool_result = self.maxpool2d(kernel_size, stride);
 
-        let dims: usize = self.dims().len();
+        let dims: usize = self.get_shape().len();
         assert!(dims >= 2, "Input tensor must have at least 2 dimensions.");
 
         let (h, w) = (self.shape[dims - 2], self.shape[dims - 1]);
@@ -1187,7 +1082,7 @@ where
 
         let padded_maxpool_tensor = Tensor {
             data: padded_maxpool_data,
-            shape: self.dims(),
+            shape: self.get_shape(),
             input_shape: vec![0],
         };
 
@@ -1219,14 +1114,23 @@ where
         let flat_index = n * (c_size * h_size * w_size) + c * (h_size * w_size) + h * w_size + w;
         self.data[flat_index]
     }
-
+}
+impl<T> Tensor<T>
+where
+    T: Copy + Clone + Send + Sync,
+    T: Copy + Default + std::ops::Mul<Output = T> + std::iter::Sum,
+    T: std::ops::Add<Output = T> + std::ops::Sub<Output = T> + std::ops::Mul<Output = T>,
+{
     pub fn conv2d(&self, kernels: &Tensor<T>, bias: &Tensor<T>, stride: usize) -> Tensor<T> {
         let (n_size, c_size, h_size, w_size) = self.get4d();
         let (k_n, k_c, k_h, k_w) = kernels.get4d();
 
-        assert!(self.dims().len() <= 4, "Supports only at most 4D input.");
         assert!(
-            kernels.dims().len() <= 4,
+            self.get_shape().len() <= 4,
+            "Supports only at most 4D input."
+        );
+        assert!(
+            kernels.get_shape().len() <= 4,
             "Supports only at most 4D filters."
         );
         // Validate shapes
@@ -1241,35 +1145,64 @@ where
         let out_w = (w_size - k_w) / stride + 1;
         let out_shape = vec![n_size, k_n, out_h, out_w];
 
-        let mut output = vec![T::default(); n_size * k_n * out_h * out_w];
-
-        for n in 0..n_size {
-            for o in 0..k_n {
-                for oh in 0..out_h {
-                    for ow in 0..out_w {
-                        let mut sum = T::default();
-
-                        // Convolution
-                        for c in 0..c_size {
-                            for kh in 0..k_h {
-                                for kw in 0..k_w {
-                                    let h = oh * stride + kh;
-                                    let w = ow * stride + kw;
-                                    sum = sum + self.get(n, c, h, w) * kernels.get(o, c, kh, kw);
-                                }
-                            }
-                        }
-
-                        // Add bias for this output channel (o)
-                        sum = sum + bias.data[o];
+        // let mut output = vec![T::default(); n_size * k_n * out_h * out_w];
+        // for n in 0..n_size {
+        //     for o in 0..k_n {
+        //         for oh in 0..out_h {
+        //             for ow in 0..out_w {
+        //                 let mut sum = T::default();
+
+        //                 // Convolution
+        //                 for c in 0..c_size {
+        //                     for kh in 0..k_h {
+        //                         for kw in 0..k_w {
+        //                             let h = oh * stride + kh;
+        //                             let w = ow * stride + kw;
+        //                             sum = sum + self.get(n, c, h, w) * kernels.get(o, c, kh, kw);
+        //                         }
+        //                     }
+        //                 }
+
+        //                 // Add bias for this output channel (o)
+        //                 sum = sum + bias.data[o];
+
+        //                 let output_index =
+        //                     n * (k_n * out_h * out_w) + o * (out_h * out_w) + oh * out_w + ow;
+        //                 output[output_index] = sum;
+        //             }
+        //         }
+        //     }
+        // }
 
-                        let output_index =
-                            n * (k_n * out_h * out_w) + o * (out_h * out_w) + oh * out_w + ow;
-                        output[output_index] = sum;
+        // Compute output in parallel
+        let output: Vec<T> = (0..n_size * k_n * out_h * out_w)
+            .into_par_iter()
+            .map(|flat_idx| {
+                // Decompose flat index into (n, o, oh, ow)
+                let n = flat_idx / (k_n * out_h * out_w);
+                let rem1 = flat_idx % (k_n * out_h * out_w);
+                let o = rem1 / (out_h * out_w);
+                let rem2 = rem1 % (out_h * out_w);
+                let oh = rem2 / out_w;
+                let ow = rem2 % out_w;
+
+                let mut sum = T::default();
+
+                // Convolution
+                for c in 0..c_size {
+                    for kh in 0..k_h {
+                        for kw in 0..k_w {
+                            let h = oh * stride + kh;
+                            let w = ow * stride + kw;
+                            sum = sum + self.get(n, c, h, w) * kernels.get(o, c, kh, kw);
+                        }
                     }
                 }
-            }
-        }
+
+                // Add bias
+                sum + bias.data[o].clone()
+            })
+            .collect();
 
         Tensor {
             data: output,
@@ -1299,7 +1232,7 @@ where
         assert!(mat_shp_pad.len() == 2 && self.shape.len() == 2);
 
         let mut new_data = vec![T::default(); mat_shp_pad.iter().product()];
-        let mat_shp_og = self.dims();
+        let mat_shp_og = self.get_shape();
         for row in 0..mat_shp_og[0] {
             for channel in 0..conv_shape_og[0] {
                 for h_in in 0..conv_shape_og[1] {
@@ -1370,43 +1303,6 @@ impl PartialEq for Tensor<Element> {
     }
 }
 
-impl Tensor<GoldilocksExt2> {
-    // /// Creates a random matrix with a given number of rows and cols.
-    // /// NOTE: doesn't take a rng as argument because to generate it in parallel it needs be sync +
-    // /// sync which is not true for basic rng core.
-    // pub fn random(shape: Vec<usize>) -> Self {
-    //     let mut rng = thread_rng();
-    //     let size = shape.iter().product();
-    //     let data = (0..size)
-    //         .map(|_| GoldilocksExt2::random(&mut rng))
-    //         .collect_vec();
-
-    //     Self { data, shape }
-    // }
-
-    /// Creates a random matrix with a given number of rows and cols.
-    /// NOTE: doesn't take a rng as argument because to generate it in parallel it needs be sync +
-    /// sync which is not true for basic rng core.
-    pub fn random_seed(shape: Vec<usize>, seed: Option<u64>) -> Self {
-        let seed = seed.unwrap_or(rand::random::<u64>()); // Use provided seed or default
-
-        let size = shape.iter().product();
-        let data = (0..size)
-            .into_par_iter()
-            .map(|i| {
-                let mut rng = StdRng::seed_from_u64(seed + i as u64);
-                GoldilocksExt2::random(&mut rng)
-            })
-            .collect::<Vec<GoldilocksExt2>>();
-
-        Self {
-            data,
-            shape,
-            input_shape: vec![0],
-        }
-    }
-}
-
 impl PartialEq for Tensor<GoldilocksExt2> {
     fn eq(&self, other: &Self) -> bool {
         self.shape == other.shape && self.data == other.data
@@ -1530,9 +1426,9 @@ mod test {
         let mat = Tensor::<Element>::random_seed(shape.clone(), Some(213));
         // println!("{}", mat);
         let new_shape = vec![shape[0].next_power_of_two(), shape[1].next_power_of_two()];
-        let new_mat = mat.pad_next_power_of_two_2d();
+        let new_mat = mat.pad_next_power_of_two();
         assert_eq!(
-            new_mat.dims(),
+            new_mat.get_shape(),
             new_shape,
             "Matrix padding to next power of two failed."
         );
@@ -1541,7 +1437,7 @@ mod test {
     impl Tensor<Element> {
         pub fn get_2d(&self, i: usize, j: usize) -> Element {
             assert!(self.is_matrix() == true);
-            self.data[i * self.dims()[1] + j]
+            self.data[i * self.get_shape()[1] + j]
         }
 
         pub fn random_eval_point(&self) -> Vec<E> {
@@ -1555,8 +1451,8 @@ mod test {
     #[test]
     fn test_tensor_mle() {
         let mat = Tensor::random(vec![3, 5]);
-        let shape = mat.dims();
-        let mat = mat.pad_next_power_of_two_2d();
+        let shape = mat.get_shape();
+        let mat = mat.pad_next_power_of_two();
         println!("matrix {}", mat);
         let mut mle = mat.clone().to_mle_2d::<E>();
         let (chosen_row, chosen_col) = (
@@ -1595,17 +1491,6 @@ mod test {
 
     type E = GoldilocksExt2;
 
-    // pub fn random_vector<E: ExtensionField>(n: usize) -> Vec<E> {
-    // let mut rng = thread_rng();
-    // let mut arr = vec![E::ZERO; n];
-    // for i in 0..n{
-    // arr[i] = E::from(i as u64);
-    // }
-    // arr
-    // (0..n).map(|_| E::from(&mut rng)).collect_vec()
-    // (0..n).map(|_| E::random(&mut rng)).collect_vec()
-    // }
-
     #[test]
     fn test_conv() {
         for i in 0..3 {
@@ -1731,9 +1616,9 @@ mod test {
         let padded = input.pad_next_power_of_two();
 
         padded
-            .dims()
+            .get_shape()
             .iter()
-            .zip(input.dims().iter())
+            .zip(input.get_shape().iter())
             .for_each(|(padded_dim, input_dim)| {
                 assert_eq!(*padded_dim, input_dim.next_power_of_two())
             });
@@ -1764,22 +1649,6 @@ mod test {
         assert_eq!(tensor_b, tensor_c);
     }
 
-    #[test]
-    fn test_tensor_pad_last_two() {
-        let shape_a = vec![3, 1, 1];
-        let tensor_a = Tensor::<Element>::new(shape_a.clone(), vec![1; shape_a.iter().product()]);
-
-        let target_dim = vec![4, 4];
-        let shape_b = vec![3, 4, 4];
-        let tensor_b = Tensor::<Element>::new(shape_b, vec![
-            1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-            0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-        ]);
-
-        let tensor_c = tensor_a.pad_last_two_dimensions(target_dim);
-        assert_eq!(tensor_b, tensor_c);
-    }
-
     #[test]
     fn test_tensor_pad_to_shape() {
         let shape_a = vec![3, 1, 1];


Reviews:
user: login: nikkolasg
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T20:25:54Z
user: login: sshravan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T20:34:54Z
user: login: sshravan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T20:55:02Z

Review comments:
diff_hunk: @@ -408,7 +408,7 @@ mod test {
     use super::compute_betas_eval;
     use crate::{
         Claim, default_transcript, pad_vector,
-        tensor::Tensor,
path: zkml/src/commit/precommit.rs
user: login: nikkolasg
body: Sorry but can we keep it singular ? Module names if we stick to singular it's just easier to remember. Most libraries I know are using singular names.
created_at: 2025-03-19T20:24:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -403,6 +394,42 @@ impl Tensor<Element> {
     pub fn filter_size(&self) -> usize {
         self.shape[2] * self.shape[2]
     }
+    /// Returns the evaluation point, in order for (row,col) addressing
+    pub fn evals_2d<F: ExtensionField>(&self) -> Vec<F> {
+        assert!(self.is_matrix(), "Tensor is not a matrix");
+        self.evals_flat()
+    }
+    pub fn evals_flat<F: ExtensionField>(&self) -> Vec<F> {
+        self.data.par_iter().map(|e| e.to_field()).collect()
+    }
+    pub fn get_conv_weights<F: ExtensionField>(&self) -> Vec<F> {
+        let mut data = vec![F::ZERO; self.data.len()];
+        for i in 0..data.len() {
+            if self.data[i] < 0 {
+                data[i] = -F::from((-self.data[i]) as u64);
+            } else {
+                data[i] = F::from(self.data[i] as u64);
path: zkml/src/tensors.rs
user: login: nikkolasg
body: let's use the `to_field()` 
created_at: 2025-03-19T20:25:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -408,7 +408,7 @@ mod test {
     use super::compute_betas_eval;
     use crate::{
         Claim, default_transcript, pad_vector,
-        tensor::Tensor,
path: zkml/src/commit/precommit.rs
user: login: sshravan
body: happy to revert 
created_at: 2025-03-19T20:34:53Z
author_association: MEMBER
diff_hunk: @@ -403,6 +394,42 @@ impl Tensor<Element> {
     pub fn filter_size(&self) -> usize {
         self.shape[2] * self.shape[2]
     }
+    /// Returns the evaluation point, in order for (row,col) addressing
+    pub fn evals_2d<F: ExtensionField>(&self) -> Vec<F> {
+        assert!(self.is_matrix(), "Tensor is not a matrix");
+        self.evals_flat()
+    }
+    pub fn evals_flat<F: ExtensionField>(&self) -> Vec<F> {
+        self.data.par_iter().map(|e| e.to_field()).collect()
+    }
+    pub fn get_conv_weights<F: ExtensionField>(&self) -> Vec<F> {
+        let mut data = vec![F::ZERO; self.data.len()];
+        for i in 0..data.len() {
+            if self.data[i] < 0 {
+                data[i] = -F::from((-self.data[i]) as u64);
+            } else {
+                data[i] = F::from(self.data[i] as u64);
path: zkml/src/tensors.rs
user: login: sshravan
body: Done [here](https://github.com/Lagrange-Labs/deep-prove/pull/77/commits/23558a7b1838f783430b01f636c752b63a2a13da#diff-7f27887ead8f7cc0b52c7adf2138ce4c708d1038aa9223ba8ae5b1aebf2c7c75R408)

created_at: 2025-03-19T20:55:02Z
author_association: MEMBER

PR comments:
user: login: nikkolasg
created_at: 2025-03-20T13:32:15Z
body: Merging to test if performance increased on the large model/dataset
author_association: CONTRIBUTOR