chore: enable rounding in FHE

src/concrete/ml/quantization/quantized_module.py

@@ -470,19 +470,6 @@ def _fhe_forward(self, *q_x: numpy.ndarray, simulate: bool = True) -> numpy.ndar
             "executing it in FHE.",
         )
 
-        # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/2888
-        # Check if rounding is being executed in FHE, which is not currently optimized
-        # Assert that there are rounding nodes in the circuit graph
-        rounding_nodes = self.fhe_circuit.graph.query_nodes(  # type: ignore[union-attr]
-            operation_filter="round_bit_pattern"
-        )
-
-        assert_true(
-            not rounding_nodes or simulate,
-            "Rounding is not currently optimized for execution in FHE. "
-            "Only simulation is allowed with a rounding operator.",
-        )
-
         results_cnp_circuit_list = []
         for i in range(q_x[0].shape[0]):
 

tests/quantization/test_quantized_module.py

@@ -343,11 +343,10 @@ def test_quantized_module_rounding_fhe(model_class, input_shape, default_configu
     # Run quantized_model in simulation
     quantized_model.forward(numpy_test, fhe="simulate")
 
-    # Try to execute the model with rounding in FHE execution mode
-    with pytest.raises(
-        AssertionError, match="Rounding is not currently optimized for execution in FHE"
-    ):
-        quantized_model.forward(numpy_test, fhe="execute")
+    # Execute the model with rounding in FHE execution mode
+    quantized_model.forward(numpy_test, fhe="execute")
+
+    # TODO: https://github.com/zama-ai/concrete-ml-internal/issues/3800
 
 
 def quantized_module_predictions_are_equal(

tests/torch/test_compile_torch.py

@@ -304,8 +304,6 @@ def accuracy_test_rounding(
     quantized_numpy_module_round_low_precision and quantized_numpy_module making sure that the
     rounding feature has the expected behavior on the model accuracy.
     """
-    # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/2888
-    assert simulate, "Rounding is not available in FHE yet."
 
     # Check that the maximum_integer_bit_width is at least 4 bits to compare the rounding
     # feature with enough precision.
@@ -393,7 +391,6 @@ def accuracy_test_rounding(
         q_x = tuple(q[[i]] for q in to_tuple(qtest))
 
         # encrypt, run, and decrypt with different precision modes
-        # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/2888
         q_result = quantized_numpy_module.quantized_forward(*q_x, fhe="simulate")
         q_result_high_precision = quantized_numpy_module_round_high_precision.quantized_forward(
             *q_x,