Merge pull request #506 from tuneinsight/505-sparse-bootstrapping

Ciphertext packing in bootstrapping

circuits/ckks/bootstrapping/bootstrapping_test.go

@@ -186,6 +186,91 @@ func TestBootstrapping(t *testing.T) {
 		})
 	})
 
+	t.Run("BootstrappingPackedWithoutRingDegreeSwitch", func(t *testing.T) {
+
+		schemeParamsLit := testPrec45
+		btpParamsLit := ParametersLiteral{}
+
+		if *flagLongTest {
+			schemeParamsLit.LogN = 16
+		}
+
+		btpParamsLit.LogN = utils.Pointy(schemeParamsLit.LogN)
+
+		params, err := ckks.NewParametersFromLiteral(schemeParamsLit)
+		require.Nil(t, err)
+
+		// Insecure params for fast testing only
+		if !*flagLongTest {
+			// Corrects the message ratio to take into account the smaller number of slots and keep the same precision
+			btpParamsLit.LogMessageRatio = utils.Pointy(DefaultLogMessageRatio + (16 - params.LogN()))
+		}
+
+		sk := rlwe.NewKeyGenerator(params).GenSecretKeyNew()
+
+		ecd := ckks.NewEncoder(params)
+		enc := rlwe.NewEncryptor(params, sk)
+		dec := rlwe.NewDecryptor(params, sk)
+
+		for _, sparsity := range []int{0, 1, 2} {
+			btpParamsLit.LogSlots = utils.Pointy(*btpParamsLit.LogN - 1 - sparsity)
+			btpParams, err := NewParametersFromLiteral(params, btpParamsLit)
+			require.Nil(t, err)
+
+			t.Logf("Params: LogN=%d/LogSlots=%d/LogQP=%f", btpParams.ResidualParameters.LogN(), btpParams.ResidualParameters.LogMaxSlots(), btpParams.ResidualParameters.LogQP())
+			t.Logf("BTPParams: LogN=%d/LogSlots=%d/LogQP=%f", btpParams.BootstrappingParameters.LogN(), btpParams.BootstrappingParameters.LogMaxSlots(), btpParams.BootstrappingParameters.LogQP())
+
+			t.Log("Generating Bootstrapping Keys for LogSlots")
+			btpKeys, _, err := btpParams.GenEvaluationKeys(sk)
+			require.Nil(t, err)
+
+			evaluator, err := NewEvaluator(btpParams, btpKeys)
+			require.Nil(t, err)
+
+			for _, slotOffset := range []int{0, 1, 2, 3} {
+				logMaxSlots := params.LogMaxSlots() - slotOffset
+				logMaxSlots = utils.Min(logMaxSlots, btpParams.LogMaxSlots())
+				values := make([]complex128, 1<<logMaxSlots)
+				for i := range values {
+					values[i] = sampling.RandComplex128(-1, 1)
+				}
+
+				values[0] = complex(0.9238795325112867, 0.3826834323650898)
+				values[1] = complex(0.9238795325112867, 0.3826834323650898)
+				if len(values) > 2 {
+					values[2] = complex(0.9238795325112867, 0.3826834323650898)
+					values[3] = complex(0.9238795325112867, 0.3826834323650898)
+				}
+
+				pt := ckks.NewPlaintext(params, 0)
+				pt.LogDimensions = ring.Dimensions{Rows: 0, Cols: logMaxSlots}
+
+				cts := make([]rlwe.Ciphertext, 11)
+				for i := range cts {
+
+					require.NoError(t, ecd.Encode(utils.RotateSlice(values, i), pt))
+
+					ct, err := enc.EncryptNew(pt)
+					require.NoError(t, err)
+
+					cts[i] = *ct
+				}
+
+				if cts, err = evaluator.BootstrapMany(cts); err != nil {
+					t.Fatal(err)
+				}
+
+				for i := range cts {
+					// Checks that the output ciphertext is at the max level of paramsN1
+					require.True(t, cts[i].Level() == params.MaxLevel())
+					require.True(t, cts[i].Scale.Equal(params.DefaultScale()))
+
+					verifyTestVectorsBootstrapping(params, ecd, dec, utils.RotateSlice(values, i), &cts[i], t)
+				}
+			}
+		}
+	})
+
 	t.Run("BootstrappingPackedWithRingDegreeSwitch", func(t *testing.T) {
 
 		schemeParamsLit := testPrec45
@@ -202,69 +287,74 @@ func TestBootstrapping(t *testing.T) {
 		params, err := ckks.NewParametersFromLiteral(schemeParamsLit)
 		require.Nil(t, err)
 
-		btpParams, err := NewParametersFromLiteral(params, btpParamsLit)
-		require.Nil(t, err)
-
 		// Insecure params for fast testing only
 		if !*flagLongTest {
-			btpParams.SlotsToCoeffsParameters.LogSlots = btpParams.BootstrappingParameters.LogN() - 1
-			btpParams.CoeffsToSlotsParameters.LogSlots = btpParams.BootstrappingParameters.LogN() - 1
-
 			// Corrects the message ratio to take into account the smaller number of slots and keep the same precision
-			btpParams.Mod1ParametersLiteral.LogMessageRatio += 16 - params.LogN()
+			btpParamsLit.LogMessageRatio = utils.Pointy(DefaultLogMessageRatio + (16 - params.LogN()))
 		}
 
-		t.Logf("Params: LogN=%d/LogSlots=%d/LogQP=%f", btpParams.ResidualParameters.LogN(), btpParams.ResidualParameters.LogMaxSlots(), btpParams.ResidualParameters.LogQP())
-		t.Logf("BTPParams: LogN=%d/LogSlots=%d/LogQP=%f", btpParams.BootstrappingParameters.LogN(), btpParams.BootstrappingParameters.LogMaxSlots(), btpParams.BootstrappingParameters.LogQP())
-
 		sk := rlwe.NewKeyGenerator(params).GenSecretKeyNew()
 
-		t.Log("Generating Bootstrapping Keys")
-		btpKeys, _, err := btpParams.GenEvaluationKeys(sk)
-		require.Nil(t, err)
-
-		evaluator, err := NewEvaluator(btpParams, btpKeys)
-		require.Nil(t, err)
-
 		ecd := ckks.NewEncoder(params)
 		enc := rlwe.NewEncryptor(params, sk)
 		dec := rlwe.NewDecryptor(params, sk)
 
-		values := make([]complex128, params.MaxSlots())
-		for i := range values {
-			values[i] = sampling.RandComplex128(-1, 1)
-		}
+		for _, sparsity := range []int{0, 1, 2} {
+			btpParamsLit.LogSlots = utils.Pointy(*btpParamsLit.LogN - 1 - sparsity)
+			btpParams, err := NewParametersFromLiteral(params, btpParamsLit)
+			require.Nil(t, err)
 
-		values[0] = complex(0.9238795325112867, 0.3826834323650898)
-		values[1] = complex(0.9238795325112867, 0.3826834323650898)
-		if len(values) > 2 {
-			values[2] = complex(0.9238795325112867, 0.3826834323650898)
-			values[3] = complex(0.9238795325112867, 0.3826834323650898)
-		}
+			t.Logf("Params: LogN=%d/LogSlots=%d/LogQP=%f", btpParams.ResidualParameters.LogN(), btpParams.ResidualParameters.LogMaxSlots(), btpParams.ResidualParameters.LogQP())
+			t.Logf("BTPParams: LogN=%d/LogSlots=%d/LogQP=%f", btpParams.BootstrappingParameters.LogN(), btpParams.BootstrappingParameters.LogMaxSlots(), btpParams.BootstrappingParameters.LogQP())
 
-		pt := ckks.NewPlaintext(params, 0)
+			t.Log("Generating Bootstrapping Keys for LogSlots")
+			btpKeys, _, err := btpParams.GenEvaluationKeys(sk)
+			require.Nil(t, err)
 
-		cts := make([]rlwe.Ciphertext, 7)
-		for i := range cts {
+			evaluator, err := NewEvaluator(btpParams, btpKeys)
+			require.Nil(t, err)
 
-			require.NoError(t, ecd.Encode(utils.RotateSlice(values, i), pt))
+			for _, slotOffset := range []int{0, 1, 2, 3} {
+				logMaxSlots := params.LogMaxSlots() - slotOffset
+				logMaxSlots = utils.Min(logMaxSlots, btpParams.LogMaxSlots())
+				values := make([]complex128, 1<<logMaxSlots)
+				for i := range values {
+					values[i] = sampling.RandComplex128(-1, 1)
+				}
 
-			ct, err := enc.EncryptNew(pt)
-			require.NoError(t, err)
+				values[0] = complex(0.9238795325112867, 0.3826834323650898)
+				values[1] = complex(0.9238795325112867, 0.3826834323650898)
+				if len(values) > 2 {
+					values[2] = complex(0.9238795325112867, 0.3826834323650898)
+					values[3] = complex(0.9238795325112867, 0.3826834323650898)
+				}
 
-			cts[i] = *ct
-		}
+				pt := ckks.NewPlaintext(params, 0)
+				pt.LogDimensions = ring.Dimensions{Rows: 0, Cols: logMaxSlots}
 
-		if cts, err = evaluator.BootstrapMany(cts); err != nil {
-			t.Fatal(err)
-		}
+				cts := make([]rlwe.Ciphertext, 11)
+				for i := range cts {
 
-		for i := range cts {
-			// Checks that the output ciphertext is at the max level of paramsN1
-			require.True(t, cts[i].Level() == params.MaxLevel())
-			require.True(t, cts[i].Scale.Equal(params.DefaultScale()))
+					require.NoError(t, ecd.Encode(utils.RotateSlice(values, i), pt))
 
-			verifyTestVectorsBootstrapping(params, ecd, dec, utils.RotateSlice(values, i), &cts[i], t)
+					ct, err := enc.EncryptNew(pt)
+					require.NoError(t, err)
+
+					cts[i] = *ct
+				}
+
+				if cts, err = evaluator.BootstrapMany(cts); err != nil {
+					t.Fatal(err)
+				}
+
+				for i := range cts {
+					// Checks that the output ciphertext is at the max level of paramsN1
+					require.True(t, cts[i].Level() == params.MaxLevel())
+					require.True(t, cts[i].Scale.Equal(params.DefaultScale()))
+
+					verifyTestVectorsBootstrapping(params, ecd, dec, utils.RotateSlice(values, i), &cts[i], t)
+				}
+			}
 		}
 	})
 

circuits/ckks/bootstrapping/evaluator.go

@@ -32,6 +32,8 @@ type Evaluator struct {
 	xPow2N1 []ring.Poly
 	// [1, x, x^2, x^4, ..., x^N2/2] / (X^N2 +1)
 	xPow2N2 []ring.Poly
+	// [1, x^-1, x^-2, x^-4, ..., x^-N1/2] / (X^N1 +1)
+	xPow2InvN1 []ring.Poly
 	// [1, x^-1, x^-2, x^-4, ..., x^-N2/2] / (X^N2 +1)
 	xPow2InvN2 []ring.Poly
 
@@ -73,9 +75,10 @@ func NewEvaluator(btpParams Parameters, evk *EvaluationKeys) (eval *Evaluator, e
 
 	if paramsN1.N() != paramsN2.N() {
 		eval.xPow2N1 = rlwe.GenXPow2NTT(paramsN1.RingQ().AtLevel(0), paramsN2.LogN(), false)
-		eval.xPow2N2 = rlwe.GenXPow2NTT(paramsN2.RingQ().AtLevel(0), paramsN2.LogN(), false)
-		eval.xPow2InvN2 = rlwe.GenXPow2NTT(paramsN2.RingQ(), paramsN2.LogN(), true)
+		eval.xPow2InvN1 = rlwe.GenXPow2NTT(paramsN1.RingQ(), paramsN1.LogN(), true)
 	}
+	eval.xPow2N2 = rlwe.GenXPow2NTT(paramsN2.RingQ().AtLevel(0), paramsN2.LogN(), false)
+	eval.xPow2InvN2 = rlwe.GenXPow2NTT(paramsN2.RingQ(), paramsN2.LogN(), true)
 
 	if btpParams.Mod1ParametersLiteral.Mod1Type == mod1.SinContinuous && btpParams.Mod1ParametersLiteral.DoubleAngle != 0 {
 		return nil, fmt.Errorf("cannot use double angle formula for Mod1Type = Sin -> must use Mod1Type = Cos")
@@ -131,11 +134,12 @@ func (eval Evaluator) ShallowCopy() *Evaluator {
 	heEvaluator := eval.Evaluator.ShallowCopy()
 
 	paramsN1 := eval.ResidualParameters
+	paramsN2 := eval.BootstrappingParameters
 
 	var DomainSwitcher ckks.DomainSwitcher
 	if paramsN1.RingType() == ring.ConjugateInvariant {
 		var err error
-		if DomainSwitcher, err = ckks.NewDomainSwitcher(eval.Parameters.BootstrappingParameters, eval.EvkCmplxToReal, eval.EvkRealToCmplx); err != nil {
+		if DomainSwitcher, err = ckks.NewDomainSwitcher(paramsN2, eval.EvkCmplxToReal, eval.EvkRealToCmplx); err != nil {
 			panic(fmt.Errorf("cannot NewBootstrapper: ckks.NewDomainSwitcher: %w", err))
 		}
 	}
@@ -150,8 +154,8 @@ func (eval Evaluator) ShallowCopy() *Evaluator {
 		xPow2N2:        eval.xPow2N2,
 		xPow2InvN2:     eval.xPow2InvN2,
 		DomainSwitcher: DomainSwitcher,
-		DFTEvaluator:   dft.NewEvaluator(paramsN1, heEvaluator),
-		Mod1Evaluator:  mod1.NewEvaluator(heEvaluator, polynomial.NewEvaluator(paramsN1, heEvaluator), eval.Mod1Parameters),
+		DFTEvaluator:   dft.NewEvaluator(paramsN2, heEvaluator),
+		Mod1Evaluator:  mod1.NewEvaluator(heEvaluator, polynomial.NewEvaluator(paramsN2, heEvaluator), eval.Mod1Parameters),
 		SkDebug:        eval.SkDebug,
 	}
 }
@@ -251,7 +255,7 @@ func (eval Evaluator) Bootstrap(ct *rlwe.Ciphertext) (*rlwe.Ciphertext, error) {
 	return &cts[0], nil
 }
 
-// BootstrapMany bootstraps a list of ciphertext and returns the list of bootstrapped ciphertexts.
+// BootstrapMany bootstraps a list of ciphertexts and returns the list of bootstrapped ciphertexts.
 func (eval Evaluator) BootstrapMany(cts []rlwe.Ciphertext) ([]rlwe.Ciphertext, error) {
 
 	var err error
@@ -283,13 +287,12 @@ func (eval Evaluator) BootstrapMany(cts []rlwe.Ciphertext) ([]rlwe.Ciphertext, e
 
 	default:
 
-		LogSlots := cts[0].LogSlots()
-		nbCiphertexts := len(cts)
-
-		if cts, err = eval.PackAndSwitchN1ToN2(cts); err != nil {
+		ctsPacked, ctxt1, ctxt2, err := eval.PackAndSwitchN1ToN2(cts)
+		if err != nil {
 			return nil, fmt.Errorf("cannot BootstrapMany: %w", err)
 		}
 
+		cts = ctsPacked
 		for i := range cts {
 			var ct *rlwe.Ciphertext
 			if ct, err = eval.Evaluate(&cts[i]); err != nil {
@@ -298,7 +301,7 @@ func (eval Evaluator) BootstrapMany(cts []rlwe.Ciphertext) ([]rlwe.Ciphertext, e
 			cts[i] = *ct
 		}
 
-		if cts, err = eval.UnpackAndSwitchN2Tn1(cts, LogSlots, nbCiphertexts); err != nil {
+		if cts, err = eval.UnpackAndSwitchN2ToN1(cts, ctxt1, ctxt2); err != nil {
 			return nil, fmt.Errorf("cannot BootstrapMany: %w", err)
 		}
 	}
@@ -820,7 +823,7 @@ func (eval Evaluator) SlotsToCoeffs(ctReal, ctImag *rlwe.Ciphertext) (ctOut *rlw
 	return eval.DFTEvaluator.SlotsToCoeffsNew(ctReal, ctImag, eval.S2CDFTMatrix)
 }
 
-func (eval Evaluator) SwitchRingDegreeN1ToN2New(ctN1 *rlwe.Ciphertext) (ctN2 *rlwe.Ciphertext) {
+func (eval Evaluator) switchRingDegreeN1ToN2New(ctN1 *rlwe.Ciphertext) (ctN2 *rlwe.Ciphertext) {
 	ctN2 = ckks.NewCiphertext(eval.BootstrappingParameters, 1, ctN1.Level())
 
 	// Sanity check, this error should never happen unless this algorithm has been improperly
@@ -831,7 +834,7 @@ func (eval Evaluator) SwitchRingDegreeN1ToN2New(ctN1 *rlwe.Ciphertext) (ctN2 *rl
 	return
 }
 
-func (eval Evaluator) SwitchRingDegreeN2ToN1New(ctN2 *rlwe.Ciphertext) (ctN1 *rlwe.Ciphertext) {
+func (eval Evaluator) switchRingDegreeN2ToN1New(ctN2 *rlwe.Ciphertext) (ctN1 *rlwe.Ciphertext) {
 	ctN1 = ckks.NewCiphertext(eval.ResidualParameters, 1, ctN2.Level())
 
 	// Sanity check, this error should never happen unless this algorithm has been improperly
@@ -864,79 +867,134 @@ func (eval Evaluator) RealToComplexNew(ctReal *rlwe.Ciphertext) (ctCmplx *rlwe.C
 	return
 }
 
-func (eval Evaluator) PackAndSwitchN1ToN2(cts []rlwe.Ciphertext) ([]rlwe.Ciphertext, error) {
+// packingContext contains the parameters used when packing (with Pack())
+type packingContext struct {
+	Params           *ckks.Parameters // Parameters of the ring we are packing to or unpacking from
+	LogMaxDimensions ring.Dimensions  // maximum dimension of a packed ciphertext (logMaxDimensions <= params.LogMaxDimensions())
+	LogSlots         int              // number of slots in a ct before packing (resp. after unpacking)
+	NbPackedCTs      int              // number of cts to be packed (resp. to be unpacked into)
+}
+
+// PackAndSwitchN1ToN2 packs the ciphertexts into N1 and switch to N2 if N1 < N2
+// then it packs the ciphertexts into N2.
+func (eval Evaluator) PackAndSwitchN1ToN2(cts []rlwe.Ciphertext) ([]rlwe.Ciphertext, *packingContext, *packingContext, error) {
 
 	var err error
+	var packN1, packN2 *packingContext
 
+	// If N1 < N2, we pack ciphertexts into N1 and then switch to N2
 	if eval.ResidualParameters.N() != eval.BootstrappingParameters.N() {
-		if cts, err = eval.Pack(cts, eval.ResidualParameters, eval.xPow2N1); err != nil {
-			return nil, fmt.Errorf("cannot PackAndSwitchN1ToN2: PackN1: %w", err)
+
+		packN1 = &packingContext{&eval.ResidualParameters, eval.ResidualParameters.LogMaxDimensions(), cts[0].LogSlots(), len(cts)}
+
+		// If the bootstrapping max slots are smaller than the max slots of N1, we only pack up to the former
+		if eval.Parameters.LogMaxSlots() < eval.ResidualParameters.LogMaxSlots() {
+			packN1.LogMaxDimensions = eval.Parameters.LogMaxDimensions()
+		}
+		if cts, err = eval.pack(cts, *packN1, eval.xPow2N1); err != nil {
+			return nil, nil, nil, fmt.Errorf("cannot PackAndSwitchN1ToN2: PackN1: %w", err)
 		}
 
 		for i := range cts {
-			cts[i] = *eval.SwitchRingDegreeN1ToN2New(&cts[i])
+			cts[i] = *eval.switchRingDegreeN1ToN2New(&cts[i])
 		}
 	}
 
-	if cts, err = eval.Pack(cts, eval.BootstrappingParameters, eval.xPow2N2); err != nil {
-		return nil, fmt.Errorf("cannot PackAndSwitchN1ToN2: PackN2: %w", err)
+	// Packing ciphertexts into N2 (up to eval.Parameters.LogMaxDimensions())
+	packN2 = &packingContext{&eval.BootstrappingParameters, eval.Parameters.LogMaxDimensions(), cts[0].LogSlots(), len(cts)}
+
+	if cts, err = eval.pack(cts, *packN2, eval.xPow2N2); err != nil {
+		return nil, nil, nil, fmt.Errorf("cannot PackAndSwitchN1ToN2: PackN1: %w", err)
 	}
 
-	return cts, nil
+	return cts, packN1, packN2, nil
 }
 
-func (eval Evaluator) UnpackAndSwitchN2Tn1(cts []rlwe.Ciphertext, LogSlots, Nb int) ([]rlwe.Ciphertext, error) {
+// UnpackAndSwitchN2ToN1 unpacks the ciphertexts into N2 and, if N1 < N2, it switches the ciphertexts
+// to N1 and unpacks further into N1
+func (eval Evaluator) UnpackAndSwitchN2ToN1(cts []rlwe.Ciphertext, ctxtN1, ctxtN2 *packingContext) ([]rlwe.Ciphertext, error) {
 
-	var err error
+	var ctsOut []rlwe.Ciphertext
 
-	if eval.ResidualParameters.N() != eval.BootstrappingParameters.N() {
-		if cts, err = eval.UnPack(cts, eval.BootstrappingParameters, LogSlots, Nb, eval.xPow2InvN2); err != nil {
+	logSlots := ctxtN2.LogSlots
+
+	// Unpack ciphertexts in N2
+	for i := range cts {
+		ctsUnpack, err := eval.unpack(&cts[i], *ctxtN2, eval.xPow2InvN2)
+
+		if err != nil {
 			return nil, fmt.Errorf("cannot UnpackAndSwitchN2Tn1: UnpackN2: %w", err)
 		}
 
-		for i := range cts {
-			cts[i] = *eval.SwitchRingDegreeN2ToN1New(&cts[i])
+		ctsOut = append(ctsOut, ctsUnpack...)
+		ctxtN2.NbPackedCTs -= len(ctsUnpack)
+	}
+
+	// If N1 != N2 (i.e. ctxtN1 != nil): 1) switch cts to N1 2) unpack the cts in N1
+	if ctxtN1 != nil {
+		var ctsN1 []rlwe.Ciphertext
+		logSlots = ctxtN1.LogSlots
+
+		for i := range ctsOut {
+			ctsOut[i] = *eval.switchRingDegreeN2ToN1New(&ctsOut[i])
 		}
+
+		for i := range ctsOut {
+			ctsUnpack, err := eval.unpack(&ctsOut[i], *ctxtN1, eval.xPow2InvN1)
+			if err != nil {
+				return nil, fmt.Errorf("cannot UnpackAndSwitchN2Tn1: UnpackN1: %w", err)
+			}
+
+			ctsN1 = append(ctsN1, ctsUnpack...)
+			ctxtN1.NbPackedCTs -= len(ctsUnpack)
+		}
+
+		ctsOut = ctsN1
 	}
 
-	for i := range cts {
-		cts[i].LogDimensions.Cols = LogSlots
+	// Set back the dimension of cts to its original value
+	for i := range ctsOut {
+		ctsOut[i].LogDimensions.Cols = logSlots
 	}
 
-	return cts, nil
+	return ctsOut, nil
 }
 
-func (eval Evaluator) UnPack(cts []rlwe.Ciphertext, params ckks.Parameters, LogSlots, Nb int, xPow2Inv []ring.Poly) ([]rlwe.Ciphertext, error) {
-	LogGap := params.LogMaxSlots() - LogSlots
+// unpack unpacks one sparse ciphertext of (log) dimension ctxt.logMaxDimensions
+// into ctxt.NbPackedCTs ciphertexts of (log) dimension {0, ctxt.LogSlots}
+func (eval Evaluator) unpack(ct *rlwe.Ciphertext, ctxt packingContext, xPow2Inv []ring.Poly) ([]rlwe.Ciphertext, error) {
+	logPackCTs := ctxt.LogMaxDimensions.Cols - ctxt.LogSlots // log of number of CTs that can be packed in one ct
 
-	if LogGap == 0 {
+	cts := []rlwe.Ciphertext{*ct}
+	if logPackCTs == 0 {
 		return cts, nil
 	}
 
-	cts = append(cts, make([]rlwe.Ciphertext, Nb-1)...)
+	n := utils.Min(ctxt.NbPackedCTs, 1<<logPackCTs) // #cts to unpack from ct
+	cts = append(cts, make([]rlwe.Ciphertext, n-1)...)
 
 	for i := 1; i < len(cts); i++ {
-		cts[i] = *cts[0].CopyNew()
+		cts[i] = *ct.CopyNew()
 	}
 
-	r := params.RingQ().AtLevel(cts[0].Level())
+	r := ctxt.Params.RingQ().AtLevel(cts[0].Level())
 
-	N := len(cts)
+	logGap := (ctxt.Params.LogMaxSlots() - ctxt.LogSlots) - 1 // log gap of CTs with params.N (minus one)
 
-	for i := 0; i < utils.Min(bits.Len64(uint64(N-1)), LogGap); i++ {
+	for i := 0; i < utils.Min(bits.Len64(uint64(n-1)), logPackCTs); i++ {
 
 		step := 1 << (i + 1)
 
-		for j := 0; j < N; j += step {
+		for j := 0; j < n; j += step {
 
 			for k := step >> 1; k < step; k++ {
 
-				if (j + k) >= N {
+				if (j + k) >= n {
 					break
 				}
 
-				r.MulCoeffsMontgomery(cts[j+k].Value[0], xPow2Inv[i], cts[j+k].Value[0])
-				r.MulCoeffsMontgomery(cts[j+k].Value[1], xPow2Inv[i], cts[j+k].Value[1])
+				r.MulCoeffsMontgomery(cts[j+k].Value[0], xPow2Inv[logGap-i], cts[j+k].Value[0])
+				r.MulCoeffsMontgomery(cts[j+k].Value[1], xPow2Inv[logGap-i], cts[j+k].Value[1])
 			}
 		}
 	}
@@ -944,28 +1002,36 @@ func (eval Evaluator) UnPack(cts []rlwe.Ciphertext, params ckks.Parameters, LogS
 	return cts, nil
 }
 
-func (eval Evaluator) Pack(cts []rlwe.Ciphertext, params ckks.Parameters, xPow2 []ring.Poly) ([]rlwe.Ciphertext, error) {
+// pack packs ctxt.NbPackedCTs sparse ciphertexts of (log) dimension {0, ctxt.LogSlots}
+// into one ciphertext of (log) dimension ctxt.logMaxDimensions
+func (eval Evaluator) pack(cts []rlwe.Ciphertext, ctxt packingContext, xPow2 []ring.Poly) ([]rlwe.Ciphertext, error) {
 
-	var LogSlots = cts[0].LogSlots()
-	RingDegree := params.N()
+	var logSlots = ctxt.LogSlots
+	var logMaxSlots = ctxt.LogMaxDimensions.Cols
+	ringDegree := ctxt.Params.N()
+
+	if logSlots > logMaxSlots {
+		return nil, fmt.Errorf("cannot Pack: cts[0].LogSlots()=%d > logMaxSlots=%d", logSlots, logMaxSlots)
+	}
 
 	for i, ct := range cts {
-		if N := ct.LogSlots(); N != LogSlots {
-			return nil, fmt.Errorf("cannot Pack: cts[%d].PlaintextLogSlots()=%d != cts[0].PlaintextLogSlots=%d", i, N, LogSlots)
+		if s := ct.LogSlots(); s != logSlots {
+			return nil, fmt.Errorf("cannot Pack: cts[%d].PlaintextLogSlots()=%d != cts[0].PlaintextLogSlots=%d", i, s, logSlots)
 		}
 
-		if N := ct.Value[0].N(); N != RingDegree {
-			return nil, fmt.Errorf("cannot Pack: cts[%d].Value[0].N()=%d != params.N()=%d", i, N, RingDegree)
+		if N := ct.Value[0].N(); N != ringDegree {
+			return nil, fmt.Errorf("cannot Pack: cts[%d].Value[0].N()=%d != params.N()=%d", i, N, ringDegree)
 		}
 	}
 
-	LogGap := params.LogMaxSlots() - LogSlots
+	logPackCTs := logMaxSlots - logSlots // log of number of CTs that can be packed in one ct
+	logGap := (ctxt.Params.LogMaxSlots() - logSlots - 1)
 
-	if LogGap == 0 {
+	if logPackCTs == 0 {
 		return cts, nil
 	}
 
-	for i := 0; i < LogGap; i++ {
+	for i := 0; i < logPackCTs; i++ {
 
 		for j := 0; j < len(cts)>>1; j++ {
 
@@ -974,10 +1040,10 @@ func (eval Evaluator) Pack(cts []rlwe.Ciphertext, params ckks.Parameters, xPow2
 
 			level := utils.Min(eve.Level(), odd.Level())
 
-			r := params.RingQ().AtLevel(level)
+			r := ctxt.Params.RingQ().AtLevel(level)
 
-			r.MulCoeffsMontgomeryThenAdd(odd.Value[0], xPow2[i], eve.Value[0])
-			r.MulCoeffsMontgomeryThenAdd(odd.Value[1], xPow2[i], eve.Value[1])
+			r.MulCoeffsMontgomeryThenAdd(odd.Value[0], xPow2[logGap-i], eve.Value[0])
+			r.MulCoeffsMontgomeryThenAdd(odd.Value[1], xPow2[logGap-i], eve.Value[1])
 
 			cts[j] = eve
 		}
@@ -990,9 +1056,8 @@ func (eval Evaluator) Pack(cts []rlwe.Ciphertext, params ckks.Parameters, xPow2
 		}
 	}
 
-	LogMaxDimensions := params.LogMaxDimensions()
 	for i := range cts {
-		cts[i].LogDimensions = LogMaxDimensions
+		cts[i].LogDimensions = ctxt.LogMaxDimensions
 	}
 
 	return cts, nil