feat: Changes on KeySwitchKey

mktfhe/bootstrap_key.go

@@ -22,9 +22,9 @@ func NewEvaluationKey[T tfhe.TorusInt](params Parameters[T]) EvaluationKey[T] {
 }
 
 // NewEvaluationKeyCustom allocates an empty EvaluationKey with custom parameters.
-func NewEvaluationKeyCustom[T tfhe.TorusInt](lweDimension, polyDegree int, bootstrapParams, keyswitchParams, relinParams tfhe.GadgetParameters[T]) EvaluationKey[T] {
+func NewEvaluationKeyCustom[T tfhe.TorusInt](lweDimension, polyDegree, glwePartialDimension int, bootstrapParams, keyswitchParams, relinParams tfhe.GadgetParameters[T]) EvaluationKey[T] {
 	return EvaluationKey[T]{
-		EvaluationKey: tfhe.NewEvaluationKeyCustom(lweDimension, 1, polyDegree, bootstrapParams, keyswitchParams),
+		EvaluationKey: tfhe.NewEvaluationKeyCustom(lweDimension, 1, polyDegree, glwePartialDimension, bootstrapParams, keyswitchParams),
 		CRSPublicKey:  tfhe.NewFourierGLevCiphertextCustom(1, polyDegree, relinParams),
 		RelinKey:      NewFourierUniEncryptionCustom(polyDegree, relinParams),
 	}

tfhe/bootstrap.go

@@ -449,17 +449,15 @@ func (e *Evaluator[T]) blindRotateOriginalAssign(ct LWECiphertext[T], lut LookUp
 }
 
 // KeySwitch switches key of ct.
-// Input ciphertext should be of length ksk.InputLWEDimension + 1,
-// and output ciphertext will be of length ksk.OutputLWEDimension + 1.
+// Input ciphertext should be of length ksk.InputLWEDimension + 1.
 func (e *Evaluator[T]) KeySwitch(ct LWECiphertext[T], ksk KeySwitchKey[T]) LWECiphertext[T] {
-	ctOut := NewLWECiphertextCustom[T](ksk.OutputLWEDimension())
+	ctOut := NewLWECiphertext(e.Parameters)
 	e.KeySwitchAssign(ct, ksk, ctOut)
 	return ctOut
 }
 
 // KeySwitchAssign switches key of ct and writes it to ctOut.
-// Input ciphertext should be of length ksk.InputLWEDimension + 1,
-// and output ciphertext should be of length ksk.OutputLWEDimension + 1.
+// Input ciphertext should be of length ksk.InputLWEDimension + 1.
 func (e *Evaluator[T]) KeySwitchAssign(ct LWECiphertext[T], ksk KeySwitchKey[T], ctOut LWECiphertext[T]) {
 	decomposed := e.decomposedBuffer(ksk.GadgetParameters)
 

tfhe/bootstrap_key.go

@@ -20,10 +20,10 @@ func NewEvaluationKey[T TorusInt](params Parameters[T]) EvaluationKey[T] {
 }
 
 // NewEvaluationKeyCustom allocates an empty EvaluationKey with custom parameters.
-func NewEvaluationKeyCustom[T TorusInt](lweDimension, glweRank, polyDegree int, bootstrapParams, keyswitchParams GadgetParameters[T]) EvaluationKey[T] {
+func NewEvaluationKeyCustom[T TorusInt](lweDimension, glweRank, polyDegree, glwePartialDimension int, bootstrapParams, keyswitchParams GadgetParameters[T]) EvaluationKey[T] {
 	return EvaluationKey[T]{
 		BootstrapKey: NewBootstrapKeyCustom(lweDimension, glweRank, polyDegree, bootstrapParams),
-		KeySwitchKey: NewKeySwitchKeyForBootstrapCustom(lweDimension, glweRank, polyDegree, keyswitchParams),
+		KeySwitchKey: NewKeySwitchKeyForBootstrapCustom(lweDimension, glwePartialDimension, keyswitchParams),
 	}
 }
 
@@ -108,7 +108,16 @@ type KeySwitchKey[T TorusInt] struct {
 }
 
 // NewKeySwitchKey allocates an empty KeySwitchingKey.
-func NewKeySwitchKey[T TorusInt](inputDimension, outputDimension int, gadgetParams GadgetParameters[T]) KeySwitchKey[T] {
+func NewKeySwitchKey[T TorusInt](params Parameters[T], inputDimension int, gadgetParams GadgetParameters[T]) KeySwitchKey[T] {
+	ksk := make([]LevCiphertext[T], inputDimension)
+	for i := 0; i < inputDimension; i++ {
+		ksk[i] = NewLevCiphertext(params, gadgetParams)
+	}
+	return KeySwitchKey[T]{Value: ksk, GadgetParameters: gadgetParams}
+}
+
+// NewKeySwitchKeyCustom allocates an empty KeySwitchingKey with custom parameters.
+func NewKeySwitchKeyCustom[T TorusInt](inputDimension, outputDimension int, gadgetParams GadgetParameters[T]) KeySwitchKey[T] {
 	ksk := make([]LevCiphertext[T], inputDimension)
 	for i := 0; i < inputDimension; i++ {
 		ksk[i] = NewLevCiphertextCustom(outputDimension, gadgetParams)
@@ -118,24 +127,19 @@ func NewKeySwitchKey[T TorusInt](inputDimension, outputDimension int, gadgetPara
 
 // NewKeySwitchKeyForBootstrap allocates an empty KeySwitchingKey for bootstrapping.
 func NewKeySwitchKeyForBootstrap[T TorusInt](params Parameters[T]) KeySwitchKey[T] {
-	return NewKeySwitchKey(params.glwePartialDimension-params.lweDimension, params.lweDimension, params.keyswitchParameters)
+	return NewKeySwitchKeyCustom(params.glwePartialDimension-params.lweDimension, params.lweDimension, params.keyswitchParameters)
 }
 
 // NewKeySwitchKeyForBootstrapCustom allocates an empty KeySwitchingKey with custom parameters.
-func NewKeySwitchKeyForBootstrapCustom[T TorusInt](lweDimension, glweRank, polyDegree int, gadgetParams GadgetParameters[T]) KeySwitchKey[T] {
-	return NewKeySwitchKey(glweRank*polyDegree-lweDimension, lweDimension, gadgetParams)
+func NewKeySwitchKeyForBootstrapCustom[T TorusInt](lweDimension, glwePartialDimension int, gadgetParams GadgetParameters[T]) KeySwitchKey[T] {
+	return NewKeySwitchKeyCustom(glwePartialDimension-lweDimension, lweDimension, gadgetParams)
 }
 
 // InputLWEDimension returns the input LWEDimension of this key.
 func (ksk KeySwitchKey[T]) InputLWEDimension() int {
 	return len(ksk.Value)
 }
 
-// OutputLWEDimension returns the output LWEDimension of this key.
-func (ksk KeySwitchKey[T]) OutputLWEDimension() int {
-	return len(ksk.Value[0].Value[0].Value) - 1
-}
-
 // Copy returns a copy of the key.
 func (ksk KeySwitchKey[T]) Copy() KeySwitchKey[T] {
 	kskCopy := make([]LevCiphertext[T], len(ksk.Value))

tfhe/bootstrap_key_marshal.go

@@ -294,7 +294,7 @@ func (ksk *KeySwitchKey[T]) ReadFrom(r io.Reader) (n int64, err error) {
 	inputDimension := int(binary.BigEndian.Uint64(metadata[16:24]))
 	outputDimension := int(binary.BigEndian.Uint64(metadata[24:32]))
 
-	*ksk = NewKeySwitchKey[T](inputDimension, outputDimension, GadgetParametersLiteral[T]{Base: T(base), Level: int(level)}.Compile())
+	*ksk = NewKeySwitchKeyCustom[T](inputDimension, outputDimension, GadgetParametersLiteral[T]{Base: T(base), Level: int(level)}.Compile())
 
 	var z T
 	switch any(z).(type) {

tfhe/bootstrap_keygen.go

@@ -108,7 +108,7 @@ func (e *Encryptor[T]) GenBootstrapKeyParallel() BootstrapKey[T] {
 // This can take a long time.
 // Use [*Encryptor.GenKeySwitchKeyParallel] for better key generation performance.
 func (e *Encryptor[T]) GenKeySwitchKey(skIn LWESecretKey[T], gadgetParams GadgetParameters[T]) KeySwitchKey[T] {
-	ksk := NewKeySwitchKey(len(skIn.Value), len(e.SecretKey.LWEKey.Value), gadgetParams)
+	ksk := NewKeySwitchKey(e.Parameters, len(skIn.Value), gadgetParams)
 
 	for i := 0; i < ksk.InputLWEDimension(); i++ {
 		for j := 0; j < gadgetParams.level; j++ {
@@ -122,7 +122,7 @@ func (e *Encryptor[T]) GenKeySwitchKey(skIn LWESecretKey[T], gadgetParams Gadget
 
 // GenKeySwitchKeyParallel samples a new keyswitch key skIn -> LWEKey in parallel.
 func (e *Encryptor[T]) GenKeySwitchKeyParallel(skIn LWESecretKey[T], gadgetParams GadgetParameters[T]) KeySwitchKey[T] {
-	ksk := NewKeySwitchKey(len(skIn.Value), len(e.SecretKey.LWEKey.Value), gadgetParams)
+	ksk := NewKeySwitchKey(e.Parameters, len(skIn.Value), gadgetParams)
 
 	workSize := ksk.InputLWEDimension() * gadgetParams.level
 	chunkCount := num.Min(runtime.NumCPU(), num.Sqrt(workSize))