Add cross-basek normalization (#90)

* added cross_basek_normalization * updated method signatures to take layouts * fixed cross-base normalization fix #91 fix #93
2026-02-10 13:16:44 +01:00 · 2025-09-30 14:40:10 +02:00
parent 4da790ea6a
commit 37e13b965c
216 changed files with 12481 additions and 7745 deletions
--- a/poulpy-core/src/automorphism/gglwe_atk.rs
+++ b/poulpy-core/src/automorphism/gglwe_atk.rs
@@ -1,43 +1,42 @@
 use poulpy_hal::{
    api::{
-        ScratchAvailable, TakeVecZnxDft, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace,
+        ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace,
        VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume,
-        VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
+        VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
    },
    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxZero},
 };

-use crate::layouts::{GGLWEAutomorphismKey, GLWECiphertext, Infos, prepared::GGLWEAutomorphismKeyPrepared};
+use crate::layouts::{GGLWEAutomorphismKey, GGLWELayoutInfos, GLWECiphertext, prepared::GGLWEAutomorphismKeyPrepared};

 impl GGLWEAutomorphismKey<Vec<u8>> {
-    #[allow(clippy::too_many_arguments)]
-    pub fn automorphism_scratch_space<B: Backend>(
+    pub fn automorphism_scratch_space<B: Backend, OUT, IN, KEY>(
        module: &Module<B>,
-        basek: usize,
-        k_out: usize,
-        k_in: usize,
-        k_ksk: usize,
-        digits: usize,
-        rank: usize,
+        out_infos: &OUT,
+        in_infos: &IN,
+        key_infos: &KEY,
    ) -> usize
    where
-        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
+        OUT: GGLWELayoutInfos,
+        IN: GGLWELayoutInfos,
+        KEY: GGLWELayoutInfos,
+        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
    {
-        GLWECiphertext::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits, rank, rank)
+        GLWECiphertext::keyswitch_scratch_space(
+            module,
+            &out_infos.glwe_layout(),
+            &in_infos.glwe_layout(),
+            key_infos,
+        )
    }

-    pub fn automorphism_inplace_scratch_space<B: Backend>(
-        module: &Module<B>,
-        basek: usize,
-        k_out: usize,
-        k_ksk: usize,
-        digits: usize,
-        rank: usize,
-    ) -> usize
+    pub fn automorphism_inplace_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
    where
-        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
+        OUT: GGLWELayoutInfos,
+        KEY: GGLWELayoutInfos,
+        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
    {
-        GGLWEAutomorphismKey::automorphism_scratch_space(module, basek, k_out, k_out, k_ksk, digits, rank)
+        GGLWEAutomorphismKey::automorphism_scratch_space(module, out_infos, out_infos, key_infos)
    }
 }

@@ -59,11 +58,15 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
            + VecZnxAutomorphism
-            + VecZnxAutomorphismInplace<B>,
-        Scratch<B>: ScratchAvailable + TakeVecZnxDft<B>,
+            + VecZnxAutomorphismInplace<B>
+            + VecZnxNormalize<B>
+            + VecZnxNormalizeTmpBytes,
+        Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
+            use crate::layouts::LWEInfos;
+
            assert_eq!(
                self.rank_in(),
                lhs.rank_in(),
@@ -93,13 +96,13 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
            )
        }

-        let cols_out: usize = rhs.rank_out() + 1;
+        let cols_out: usize = (rhs.rank_out() + 1).into();

        let p: i64 = lhs.p();
-        let p_inv = module.galois_element_inv(p);
+        let p_inv: i64 = module.galois_element_inv(p);

-        (0..self.rank_in()).for_each(|col_i| {
-            (0..self.rows()).for_each(|row_j| {
+        (0..self.rank_in().into()).for_each(|col_i| {
+            (0..self.rows().into()).for_each(|row_j| {
                let mut res_ct: GLWECiphertext<&mut [u8]> = self.at_mut(row_j, col_i);
                let lhs_ct: GLWECiphertext<&[u8]> = lhs.at(row_j, col_i);

@@ -118,8 +121,8 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
            });
        });

-        (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
-            (0..self.rank_in()).for_each(|col_j| {
+        (self.rows().min(lhs.rows()).into()..self.rows().into()).for_each(|row_i| {
+            (0..self.rank_in().into()).for_each(|col_j| {
                self.at_mut(row_i, col_j).data.zero();
            });
        });
@@ -143,8 +146,10 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
            + VecZnxAutomorphism
-            + VecZnxAutomorphismInplace<B>,
-        Scratch<B>: ScratchAvailable + TakeVecZnxDft<B>,
+            + VecZnxAutomorphismInplace<B>
+            + VecZnxNormalize<B>
+            + VecZnxNormalizeTmpBytes,
+        Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
@@ -164,13 +169,13 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
            );
        }

-        let cols_out: usize = rhs.rank_out() + 1;
+        let cols_out: usize = (rhs.rank_out() + 1).into();

        let p: i64 = self.p();
        let p_inv = module.galois_element_inv(p);

-        (0..self.rank_in()).for_each(|col_i| {
-            (0..self.rows()).for_each(|row_j| {
+        (0..self.rank_in().into()).for_each(|col_i| {
+            (0..self.rows().into()).for_each(|row_j| {
                let mut res_ct: GLWECiphertext<&mut [u8]> = self.at_mut(row_j, col_i);

                // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
--- a/poulpy-core/src/automorphism/ggsw_ct.rs
+++ b/poulpy-core/src/automorphism/ggsw_ct.rs
@@ -1,67 +1,66 @@
 use poulpy_hal::{
    api::{
-        ScratchAvailable, TakeVecZnxBig, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes,
-        VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxDftCopy,
-        VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd,
-        VmpApplyDftToDftTmpBytes,
+        ScratchAvailable, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace,
+        VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace, VecZnxDftAllocBytes,
+        VecZnxDftApply, VecZnxDftCopy, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeTmpBytes,
+        VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
    },
    layouts::{Backend, DataMut, DataRef, Module, Scratch},
 };

 use crate::layouts::{
-    GGSWCiphertext, GLWECiphertext, Infos,
+    GGLWELayoutInfos, GGSWCiphertext, GGSWInfos, GLWECiphertext,
    prepared::{GGLWEAutomorphismKeyPrepared, GGLWETensorKeyPrepared},
 };

 impl GGSWCiphertext<Vec<u8>> {
-    #[allow(clippy::too_many_arguments)]
-    pub fn automorphism_scratch_space<B: Backend>(
+    pub fn automorphism_scratch_space<B: Backend, OUT, IN, KEY, TSK>(
        module: &Module<B>,
-        basek: usize,
-        k_out: usize,
-        k_in: usize,
-        k_ksk: usize,
-        digits_ksk: usize,
-        k_tsk: usize,
-        digits_tsk: usize,
-        rank: usize,
+        out_infos: &OUT,
+        in_infos: &IN,
+        key_infos: &KEY,
+        tsk_infos: &TSK,
    ) -> usize
    where
+        OUT: GGSWInfos,
+        IN: GGSWInfos,
+        KEY: GGLWELayoutInfos,
+        TSK: GGLWELayoutInfos,
        Module<B>: VecZnxDftAllocBytes
            + VmpApplyDftToDftTmpBytes
            + VecZnxBigAllocBytes
            + VecZnxNormalizeTmpBytes
            + VecZnxBigNormalizeTmpBytes,
    {
-        let out_size: usize = k_out.div_ceil(basek);
-        let ci_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, out_size);
-        let ks_internal: usize =
-            GLWECiphertext::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits_ksk, rank, rank);
-        let expand: usize = GGSWCiphertext::expand_row_scratch_space(module, basek, k_out, k_tsk, digits_tsk, rank);
+        let out_size: usize = out_infos.size();
+        let ci_dft: usize = module.vec_znx_dft_alloc_bytes((key_infos.rank_out() + 1).into(), out_size);
+        let ks_internal: usize = GLWECiphertext::keyswitch_scratch_space(
+            module,
+            &out_infos.glwe_layout(),
+            &in_infos.glwe_layout(),
+            key_infos,
+        );
+        let expand: usize = GGSWCiphertext::expand_row_scratch_space(module, out_infos, tsk_infos);
        ci_dft + (ks_internal | expand)
    }

-    #[allow(clippy::too_many_arguments)]
-    pub fn automorphism_inplace_scratch_space<B: Backend>(
+    pub fn automorphism_inplace_scratch_space<B: Backend, OUT, KEY, TSK>(
        module: &Module<B>,
-        basek: usize,
-        k_out: usize,
-        k_ksk: usize,
-        digits_ksk: usize,
-        k_tsk: usize,
-        digits_tsk: usize,
-        rank: usize,
+        out_infos: &OUT,
+        key_infos: &KEY,
+        tsk_infos: &TSK,
    ) -> usize
    where
+        OUT: GGSWInfos,
+        KEY: GGLWELayoutInfos,
+        TSK: GGLWELayoutInfos,
        Module<B>: VecZnxDftAllocBytes
            + VmpApplyDftToDftTmpBytes
            + VecZnxBigAllocBytes
            + VecZnxNormalizeTmpBytes
            + VecZnxBigNormalizeTmpBytes,
    {
-        GGSWCiphertext::automorphism_scratch_space(
-            module, basek, k_out, k_out, k_ksk, digits_ksk, k_tsk, digits_tsk, rank,
-        )
+        GGSWCiphertext::automorphism_scratch_space(module, out_infos, out_infos, key_infos, tsk_infos)
    }
 }

@@ -88,13 +87,18 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
            + VecZnxNormalizeTmpBytes
            + VecZnxDftCopy<B>
            + VecZnxDftAddInplace<B>
-            + VecZnxIdftApplyTmpA<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnxBig<B>,
+            + VecZnxIdftApplyTmpA<B>
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnxBig<B> + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
-            assert_eq!(self.n(), auto_key.n());
-            assert_eq!(lhs.n(), auto_key.n());
+            use crate::layouts::{GLWEInfos, LWEInfos};
+
+            assert_eq!(self.n(), module.n() as u32);
+            assert_eq!(lhs.n(), module.n() as u32);
+            assert_eq!(auto_key.n(), module.n() as u32);
+            assert_eq!(tensor_key.n(), module.n() as u32);

            assert_eq!(
                self.rank(),
@@ -105,36 +109,23 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
            );
            assert_eq!(
                self.rank(),
-                auto_key.rank(),
+                auto_key.rank_out(),
                "ggsw_in rank: {} != auto_key rank: {}",
                self.rank(),
-                auto_key.rank()
+                auto_key.rank_out()
            );
            assert_eq!(
                self.rank(),
-                tensor_key.rank(),
+                tensor_key.rank_out(),
                "ggsw_in rank: {} != tensor_key rank: {}",
                self.rank(),
-                tensor_key.rank()
+                tensor_key.rank_out()
            );
-            assert!(
-                scratch.available()
-                    >= GGSWCiphertext::automorphism_scratch_space(
-                        module,
-                        self.basek(),
-                        self.k(),
-                        lhs.k(),
-                        auto_key.k(),
-                        auto_key.digits(),
-                        tensor_key.k(),
-                        tensor_key.digits(),
-                        self.rank(),
-                    )
-            )
+            assert!(scratch.available() >= GGSWCiphertext::automorphism_scratch_space(module, self, lhs, auto_key, tensor_key))
        };

        // Keyswitch the j-th row of the col 0
-        (0..lhs.rows()).for_each(|row_i| {
+        (0..lhs.rows().into()).for_each(|row_i| {
            // Key-switch column 0, i.e.
            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
            self.at_mut(row_i, 0)
@@ -164,11 +155,12 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
            + VecZnxNormalizeTmpBytes
            + VecZnxDftCopy<B>
            + VecZnxDftAddInplace<B>
-            + VecZnxIdftApplyTmpA<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnxBig<B>,
+            + VecZnxIdftApplyTmpA<B>
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnxBig<B> + TakeVecZnx,
    {
        // Keyswitch the j-th row of the col 0
-        (0..self.rows()).for_each(|row_i| {
+        (0..self.rows().into()).for_each(|row_i| {
            // Key-switch column 0, i.e.
            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
            self.at_mut(row_i, 0)
--- a/poulpy-core/src/automorphism/glwe_ct.rs
+++ b/poulpy-core/src/automorphism/glwe_ct.rs
@@ -1,44 +1,38 @@
 use poulpy_hal::{
    api::{
-        ScratchAvailable, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace,
-        VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallAInplace, VecZnxBigSubSmallBInplace,
-        VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VmpApplyDftToDft, VmpApplyDftToDftAdd,
-        VmpApplyDftToDftTmpBytes,
+        ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace,
+        VecZnxBigAutomorphismInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallInplace,
+        VecZnxBigSubSmallNegateInplace, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxNormalize,
+        VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
    },
    layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnxBig},
 };

-use crate::layouts::{GLWECiphertext, Infos, prepared::GGLWEAutomorphismKeyPrepared};
+use crate::layouts::{GGLWELayoutInfos, GLWECiphertext, GLWEInfos, LWEInfos, prepared::GGLWEAutomorphismKeyPrepared};

 impl GLWECiphertext<Vec<u8>> {
-    #[allow(clippy::too_many_arguments)]
-    pub fn automorphism_scratch_space<B: Backend>(
+    pub fn automorphism_scratch_space<B: Backend, OUT, IN, KEY>(
        module: &Module<B>,
-        basek: usize,
-        k_out: usize,
-        k_in: usize,
-        k_ksk: usize,
-        digits: usize,
-        rank: usize,
+        out_infos: &OUT,
+        in_infos: &IN,
+        key_infos: &KEY,
    ) -> usize
    where
-        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
+        OUT: GLWEInfos,
+        IN: GLWEInfos,
+        KEY: GGLWELayoutInfos,
+        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
    {
-        Self::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits, rank, rank)
+        Self::keyswitch_scratch_space(module, out_infos, in_infos, key_infos)
    }

-    pub fn automorphism_inplace_scratch_space<B: Backend>(
-        module: &Module<B>,
-        basek: usize,
-        k_out: usize,
-        k_ksk: usize,
-        digits: usize,
-        rank: usize,
-    ) -> usize
+    pub fn automorphism_inplace_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
    where
-        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
+        OUT: GLWEInfos,
+        KEY: GGLWELayoutInfos,
+        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
    {
-        Self::keyswitch_inplace_scratch_space(module, basek, k_out, k_ksk, digits, rank)
+        Self::keyswitch_inplace_scratch_space(module, out_infos, key_infos)
    }
 }

@@ -59,11 +53,13 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxIdftApplyConsume<B>
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
-            + VecZnxAutomorphismInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxAutomorphismInplace<B>
+            + VecZnxNormalize<B>
+            + VecZnxNormalizeTmpBytes,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        self.keyswitch(module, lhs, &rhs.key, scratch);
-        (0..self.rank() + 1).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i, scratch);
        })
    }
@@ -83,11 +79,13 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxIdftApplyConsume<B>
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
-            + VecZnxAutomorphismInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxAutomorphismInplace<B>
+            + VecZnxNormalize<B>
+            + VecZnxNormalizeTmpBytes,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        self.keyswitch_inplace(module, &rhs.key, scratch);
-        (0..self.rank() + 1).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i, scratch);
        })
    }
@@ -108,19 +106,29 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxIdftApplyConsume<B>
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
-            + VecZnxBigAutomorphismInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigAutomorphismInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch(module, lhs, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
            module.vec_znx_big_add_small_inplace(&mut res_big, i, &lhs.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }

@@ -139,19 +147,29 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxIdftApplyConsume<B>
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
-            + VecZnxBigAutomorphismInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigAutomorphismInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch_inplace(module, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = self.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
            module.vec_znx_big_add_small_inplace(&mut res_big, i, &self.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }

@@ -172,23 +190,33 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
            + VecZnxBigAutomorphismInplace<B>
-            + VecZnxBigSubSmallAInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigSubSmallInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch(module, lhs, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
-            module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, &lhs.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_sub_small_inplace(&mut res_big, i, &lhs.data, i);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }

-    pub fn automorphism_sub_ab_inplace<DataRhs: DataRef, B: Backend>(
+    pub fn automorphism_sub_inplace<DataRhs: DataRef, B: Backend>(
        &mut self,
        module: &Module<B>,
        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
@@ -204,23 +232,33 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
            + VecZnxBigAutomorphismInplace<B>
-            + VecZnxBigSubSmallAInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigSubSmallInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch_inplace(module, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = self.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
-            module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, &self.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_sub_small_inplace(&mut res_big, i, &self.data, i);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }

-    pub fn automorphism_sub_ba<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
+    pub fn automorphism_sub_negate<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
        &mut self,
        module: &Module<B>,
        lhs: &GLWECiphertext<DataLhs>,
@@ -237,23 +275,33 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
            + VecZnxBigAutomorphismInplace<B>
-            + VecZnxBigSubSmallBInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigSubSmallNegateInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch(module, lhs, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
-            module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, &lhs.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, &lhs.data, i);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }

-    pub fn automorphism_sub_ba_inplace<DataRhs: DataRef, B: Backend>(
+    pub fn automorphism_sub_negate_inplace<DataRhs: DataRef, B: Backend>(
        &mut self,
        module: &Module<B>,
        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
@@ -269,19 +317,29 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
            + VecZnxBigAutomorphismInplace<B>
-            + VecZnxBigSubSmallBInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigSubSmallNegateInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch_inplace(module, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = self.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
-            module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, &self.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, &self.data, i);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }
 }