Added tmp_mat_znx_dft to Scratch

backend/src/lib.rs

@@ -248,4 +248,22 @@ impl Scratch {
             Self::new(rem_slice),
         )
     }
+
+    pub fn tmp_mat_znx_dft<B: Backend>(
+        &mut self,
+        module: &Module<B>,
+        rows: usize,
+        cols_in: usize,
+        cols_out: usize,
+        size: usize,
+    ) -> (MatZnxDft<&mut [u8], B>, &mut Self) {
+        let (take_slice, rem_slice) = Self::take_slice_aligned(
+            &mut self.data,
+            module.bytes_of_mat_znx_dft(rows, cols_in, cols_out, size),
+        );
+        (
+            MatZnxDft::from_data(take_slice, module.n(), rows, cols_in, cols_out, size),
+            Self::new(rem_slice),
+        )
+    }
 }

backend/src/mat_znx_dft.rs

@@ -198,3 +198,17 @@ where
         }
     }
 }
+
+impl<D, B: Backend> MatZnxDft<D, B> {
+    pub(crate) fn from_data(data: D, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
+        Self {
+            data,
+            n,
+            rows,
+            cols_in,
+            cols_out,
+            size,
+            _phantom: PhantomData,
+        }
+    }
+}