fix(gpu): fix wrong number of blocks used in cast

zama-ai · Feb 21, 2025 · bccc919 · bccc919
1 parent 5de1445
commit bccc919
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/tfhe/src/integer/gpu/server_key/radix/mod.rs b/tfhe/src/integer/gpu/server_key/radix/mod.rs
@@ -524,7 +524,7 @@ impl CudaServerKey {
         let extended_ct_info = ct
             .as_ref()
             .info
-            .after_extend_radix_with_trivial_zero_blocks_lsb(num_blocks);
+            .after_extend_radix_with_trivial_zero_blocks_lsb(new_num_blocks);
         T::from(CudaRadixCiphertext::new(extended_ct_list, extended_ct_info))
     }
 
@@ -608,7 +608,7 @@ impl CudaServerKey {
         let extended_ct_info = ct
             .as_ref()
             .info
-            .after_extend_radix_with_trivial_zero_blocks_msb(num_blocks);
+            .after_extend_radix_with_trivial_zero_blocks_msb(new_num_blocks);
         T::from(CudaRadixCiphertext::new(extended_ct_list, extended_ct_info))
     }
 
@@ -692,7 +692,7 @@ impl CudaServerKey {
             ciphertext_modulus,
         );
 
-        let trimmed_ct_info = ct.as_ref().info.after_trim_radix_blocks_lsb(num_blocks);
+        let trimmed_ct_info = ct.as_ref().info.after_trim_radix_blocks_lsb(new_num_blocks);
         T::from(CudaRadixCiphertext::new(trimmed_ct_list, trimmed_ct_info))
     }
 
@@ -775,8 +775,7 @@ impl CudaServerKey {
             LweCiphertextCount(new_num_blocks),
             ciphertext_modulus,
         );
-
-        let trimmed_ct_info = ct.as_ref().info.after_trim_radix_blocks_msb(num_blocks);
+        let trimmed_ct_info = ct.as_ref().info.after_trim_radix_blocks_msb(new_num_blocks);
         T::from(CudaRadixCiphertext::new(trimmed_ct_list, trimmed_ct_info))
     }