Feat/integrate fhe sum and quantized dict for input and leaves

src/concrete/ml/quantization/__init__.py

@@ -1,6 +1,12 @@
 """Modules for quantization."""
 from .base_quantized_op import QuantizedOp
-from .post_training import PostTrainingAffineQuantization, PostTrainingQATImporter, get_n_bits_dict
+from .post_training import (
+    PostTrainingAffineQuantization,
+    PostTrainingQATImporter,
+    _get_n_bits_dict_trees,
+    _inspect_tree_n_bits,
+    get_n_bits_dict,
+)
 from .quantized_module import QuantizedModule
 from .quantized_ops import (
     QuantizedAbs,

src/concrete/ml/quantization/post_training.py

@@ -23,6 +23,99 @@
 from .quantized_ops import QuantizedBrevitasQuant
 from .quantizers import QuantizationOptions, QuantizedArray, UniformQuantizer
 
+# pylint: disable=too-many-lines
+
+
+def _inspect_tree_n_bits(n_bits):
+    """Validate the 'n_bits' parameter for tree-based models.
+
+    This function checks whether 'n_bits' is a valid integer or dictionary.
+    - If 'n_bits' is an integer, it must be a non-null positive, its value is assigned to
+        'op_inputs' and 'op_leaves' bits
+    - If it is a dictionary, it should contain integer values for keys 'op_leaves' and 'op_inputs',
+        where 'op_leaves' should not exceed 'op_inputs'.
+
+    The function raises a ValueError with a descriptive message if 'n_bits' does not meet
+    these criteria.
+
+    Args:
+        n_bits (int, Dict[str, int]): number of bits for quantization, can be a single value or
+            a dictionary with the following keys :
+            - "op_inputs" (mandatory): number of bits to quantize the input values
+            - "op_leaves" (optional): number of bits to quantize the leaves, must be less than or
+                equal to 'op_inputs. defaults to the value of 'op_inputs if not specified.
+
+    Raises:
+        ValueError: If 'n_bits' does not conform to the required format or value constraints.
+    """
+
+    detailed_message = (
+        "Invalid 'n_bits', either pass a strictly positive integer or a dictionary containing "
+        "integer values for the following keys:\n"
+        "- 'op_inputs' (mandatory): number of bits to quantize the input values\n"
+        "- 'op_leaves' (optional): number of bits to quantize the leaves, must be less than or "
+        "equal to 'op_inputs'. Defaults to the value of 'op_inputs' if not specified."
+        "When using a single integer for n_bits, its value is assigned to 'op_inputs' and "
+        "'op_leaves' bits.\n"
+    )
+
+    error_message = ""
+
+    if isinstance(n_bits, int):
+        if n_bits <= 0:
+            error_message = "n_bits must be a strictly positive integer"
+    elif isinstance(n_bits, dict):
+        if "op_inputs" not in n_bits.keys():
+            error_message = "Invalid keys in `n_bits` dictionary. The key 'op_inputs' is mandatory"
+        elif set(n_bits.keys()) - {"op_leaves", "op_inputs"}:
+            error_message = (
+                "Invalid keys in 'n_bits' dictionary. Only 'op_inputs' (mandatory) and 'op_leaves' "
+                "(optional) are allowed"
+            )
+        elif not all(isinstance(value, int) and value > 0 for value in n_bits.values()):
+            error_message = "All values in 'n_bits' dictionary must be strictly positive integers"
+
+        elif n_bits.get("op_leaves", 0) > n_bits.get("op_inputs", 0):
+            error_message = "'op_leaves' must be less than or equal to 'op_inputs'"
+    else:
+        error_message = "n_bits must be either an integer or a dictionary"
+
+    if len(error_message) > 0:
+        raise ValueError(
+            f"{error_message}. Got '{type(n_bits)}' and '{n_bits}' value.\n{detailed_message}"
+        )
+
+
+def _get_n_bits_dict_trees(n_bits: Union[int, Dict[str, int]]) -> Dict[str, int]:
+    """Convert the n_bits parameter into a proper dictionary for tree based-models.
+
+    Args:
+        n_bits (int, Dict[str, int]): number of bits for quantization, can be a single value or
+            a dictionary with the following keys :
+            - "op_inputs" (mandatory): number of bits to quantize the input values
+            - "op_leaves" (optional): number of bits to quantize the leaves, must be less than or
+                equal to 'op_inputs'. defaults to the value of "op_inputs" if not specified.
+
+            When using a single integer for n_bits, its value is assigned to "op_inputs" and
+            "op_leaves" bits.
+
+    Returns:
+        n_bits_dict (Dict[str, int]): A dictionary properly representing the number of bits to use
+            for quantization.
+    """
+
+    _inspect_tree_n_bits(n_bits)
+
+    # If a single integer is passed, we use a default value for the model's input and leaves
+    if isinstance(n_bits, int):
+        return {"op_inputs": n_bits, "op_leaves": n_bits}
+
+    # Default 'op_leaves' to 'op_inputs' if not specified
+    if "op_leaves" not in n_bits:
+        n_bits["op_leaves"] = n_bits["op_inputs"]
+
+    return n_bits
+
 
 def get_n_bits_dict(n_bits: Union[int, Dict[str, int]]) -> Dict[str, int]:
     """Convert the n_bits parameter into a proper dictionary.

src/concrete/ml/sklearn/base.py

@@ -49,7 +49,13 @@
 # The sigmoid and softmax functions are already defined in the ONNX module and thus are imported
 # here in order to avoid duplicating them.
 from ..onnx.ops_impl import numpy_sigmoid, numpy_softmax
-from ..quantization import PostTrainingQATImporter, QuantizedArray, get_n_bits_dict
+from ..quantization import (
+    PostTrainingQATImporter,
+    QuantizedArray,
+    _get_n_bits_dict_trees,
+    _inspect_tree_n_bits,
+    get_n_bits_dict,
+)
 from ..quantization.quantized_module import QuantizedModule, _get_inputset_generator
 from ..quantization.quantizers import (
     QuantizationOptions,
@@ -96,7 +102,7 @@
 # Enable rounding feature for all tree-based models by default
 # Note: This setting is fixed and cannot be altered by users
 # However, for internal testing purposes, we retain the capability to disable this feature
-os.environ["TREES_USE_ROUNDING"] = "1"
+os.environ["TREES_USE_ROUNDING"] = os.environ.get("TREES_USE_ROUNDING", "1")
 
 # pylint: disable=too-many-public-methods
 
@@ -1281,17 +1287,32 @@ def __init_subclass__(cls):
                 _TREE_MODELS.add(cls)
                 _ALL_SKLEARN_MODELS.add(cls)
 
-    def __init__(self, n_bits: int):
+    def __init__(self, n_bits: Union[int, Dict[str, int]]):
         """Initialize the TreeBasedEstimatorMixin.
 
         Args:
-            n_bits (int): The number of bits used for quantization.
+            n_bits (int, Dict[str, int]): Number of bits to quantize the model. If an int is passed
+                for n_bits, the value will be used for quantizing inputs and leaves. If a dict is
+                passed, then it should contain "op_inputs" and "op_leaves" as keys with
+                corresponding number of quantization bits so that:
+                    - op_inputs (mandatory): number of bits to quantize the input values
+                    - op_leaves (optional): number of bits to quantize the leaves
+                Default to 6.
         """
-        self.n_bits: int = n_bits
+
+        # Check if 'n_bits' is a valid value.
+        _inspect_tree_n_bits(n_bits)
+
+        self.n_bits: Union[int, Dict[str, int]] = n_bits
 
         #: The model's inference function. Is None if the model is not fitted.
         self._tree_inference: Optional[Callable] = None
 
+        #: Wether to perform the sum of the output's tree ensembles in FHE or not.
+        # By default, the decision of the tree ensembles is made in clear (not in FHE).
+        # This attribute should not be modified by users.
+        self._fhe_ensembling = False
+
         BaseEstimator.__init__(self)
 
     def fit(self, X: Data, y: Target, **fit_parameters):
@@ -1304,9 +1325,14 @@ def fit(self, X: Data, y: Target, **fit_parameters):
 
         q_X = numpy.zeros_like(X)
 
+        # Convert the n_bits attribute into a proper dictionary
+        self.n_bits = _get_n_bits_dict_trees(self.n_bits)
+
         # Quantization of each feature in X
         for i in range(X.shape[1]):
-            input_quantizer = QuantizedArray(n_bits=self.n_bits, values=X[:, i]).quantizer
+            input_quantizer = QuantizedArray(
+                n_bits=self.n_bits["op_inputs"], values=X[:, i]
+            ).quantizer
             self.input_quantizers.append(input_quantizer)
             q_X[:, i] = input_quantizer.quant(X[:, i])
 
@@ -1319,7 +1345,7 @@ def fit(self, X: Data, y: Target, **fit_parameters):
         # Check that the underlying sklearn model has been set and fit
         assert self.sklearn_model is not None, self._sklearn_model_is_not_fitted_error_message()
 
-        # Convert the tree inference with Numpy operators
+        # Enable rounding feature
         enable_rounding = os.environ.get("TREES_USE_ROUNDING", "1") == "1"
 
         if not enable_rounding:
@@ -1332,12 +1358,14 @@ def fit(self, X: Data, y: Target, **fit_parameters):
                 stacklevel=2,
             )
 
+        # Convert the tree inference with Numpy operators
         self._tree_inference, self.output_quantizers, self.onnx_model_ = tree_to_numpy(
             self.sklearn_model,
             q_X,
             use_rounding=enable_rounding,
+            fhe_ensembling=self._fhe_ensembling,
             framework=self.framework,
-            output_n_bits=self.n_bits,
+            output_n_bits=self.n_bits["op_leaves"],
         )
 
         self._is_fitted = True
@@ -1412,10 +1440,13 @@ def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray:
         # Sum all tree outputs
         # Remove the sum once we handle multi-precision circuits
         # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/451
-        y_preds = numpy.sum(y_preds, axis=-1)
+        if not self._fhe_ensembling:
+            y_preds = numpy.sum(y_preds, axis=-1)
 
-        assert_true(y_preds.ndim == 2, "y_preds should be a 2D array")
-        return y_preds
+            assert_true(y_preds.ndim == 2, "y_preds should be a 2D array")
+            return y_preds
+
+        return super().post_processing(y_preds)
 
 
 class BaseTreeRegressorMixin(BaseTreeEstimatorMixin, sklearn.base.RegressorMixin, ABC):
@@ -1841,12 +1872,15 @@ def __init__(self, n_bits: int = 3):
                 quantizing inputs and X_fit. Default to 3.
         """
         self.n_bits: int = n_bits
+
         # _q_fit_X: In distance metric algorithms, `_q_fit_X` stores the training set to compute
         # the similarity or distance measures. There is no `weights` attribute because there isn't
         # a training phase
         self._q_fit_X: numpy.ndarray
+
         # _y: Labels of `_q_fit_X`
         self._y: numpy.ndarray
+
         # _q_fit_X_quantizer: The quantizer to use for quantizing the model's training set
         self._q_fit_X_quantizer: Optional[UniformQuantizer] = None
 

src/concrete/ml/sklearn/rf.py

@@ -1,5 +1,5 @@
 """Implement RandomForest models."""
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 import numpy
 import sklearn.ensemble
@@ -19,7 +19,7 @@ class RandomForestClassifier(BaseTreeClassifierMixin):
     # pylint: disable-next=too-many-arguments
     def __init__(
         self,
-        n_bits: int = 6,
+        n_bits: Union[int, Dict[str, int]] = 6,
         n_estimators=20,
         criterion="gini",
         max_depth=4,
@@ -84,6 +84,7 @@ def dump_dict(self) -> Dict[str, Any]:
         metadata["onnx_model_"] = self.onnx_model_
         metadata["framework"] = self.framework
         metadata["post_processing_params"] = self.post_processing_params
+        metadata["_fhe_ensembling"] = self._fhe_ensembling
 
         # Scikit-Learn
         metadata["n_estimators"] = self.n_estimators
@@ -120,11 +121,13 @@ def load_dict(cls, metadata: Dict):
         obj.framework = metadata["framework"]
         obj.onnx_model_ = metadata["onnx_model_"]
         obj.output_quantizers = metadata["output_quantizers"]
+        obj._fhe_ensembling = metadata["_fhe_ensembling"]
         obj._tree_inference = tree_to_numpy(
             obj.sklearn_model,
             numpy.zeros((len(obj.input_quantizers),))[None, ...],
             framework=obj.framework,
-            output_n_bits=obj.n_bits,
+            output_n_bits=obj.n_bits["op_leaves"] if isinstance(obj.n_bits, Dict) else obj.n_bits,
+            fhe_ensembling=obj._fhe_ensembling,
         )[0]
         obj.post_processing_params = metadata["post_processing_params"]
 
@@ -162,7 +165,7 @@ class RandomForestRegressor(BaseTreeRegressorMixin):
     # pylint: disable-next=too-many-arguments
     def __init__(
         self,
-        n_bits: int = 6,
+        n_bits: Union[int, Dict[str, int]] = 6,
         n_estimators=20,
         criterion="squared_error",
         max_depth=4,
@@ -219,6 +222,7 @@ def dump_dict(self) -> Dict[str, Any]:
         metadata["onnx_model_"] = self.onnx_model_
         metadata["framework"] = self.framework
         metadata["post_processing_params"] = self.post_processing_params
+        metadata["_fhe_ensembling"] = self._fhe_ensembling
 
         # Scikit-Learn
         metadata["n_estimators"] = self.n_estimators
@@ -255,11 +259,13 @@ def load_dict(cls, metadata: Dict):
         obj.framework = metadata["framework"]
         obj.onnx_model_ = metadata["onnx_model_"]
         obj.output_quantizers = metadata["output_quantizers"]
+        obj._fhe_ensembling = metadata["_fhe_ensembling"]
         obj._tree_inference = tree_to_numpy(
             obj.sklearn_model,
             numpy.zeros((len(obj.input_quantizers),))[None, ...],
             framework=obj.framework,
-            output_n_bits=obj.n_bits,
+            output_n_bits=obj.n_bits["op_leaves"] if isinstance(obj.n_bits, Dict) else obj.n_bits,
+            fhe_ensembling=obj._fhe_ensembling,
         )[0]
         obj.post_processing_params = metadata["post_processing_params"]
 

src/concrete/ml/sklearn/tree.py

@@ -1,5 +1,5 @@
 """Implement DecisionTree models."""
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 import numpy
 import sklearn.tree
@@ -31,7 +31,7 @@ def __init__(
         min_impurity_decrease=0.0,
         class_weight=None,
         ccp_alpha: float = 0.0,
-        n_bits: int = 6,
+        n_bits: Union[int, Dict[str, int]] = 6,
     ):
         """Initialize the DecisionTreeClassifier.
 
@@ -84,6 +84,7 @@ def dump_dict(self) -> Dict[str, Any]:
         metadata["onnx_model_"] = self.onnx_model_
         metadata["framework"] = self.framework
         metadata["post_processing_params"] = self.post_processing_params
+        metadata["_fhe_ensembling"] = self._fhe_ensembling
 
         # Scikit-Learn
         metadata["criterion"] = self.criterion
@@ -115,11 +116,13 @@ def load_dict(cls, metadata: Dict):
         obj.framework = metadata["framework"]
         obj.onnx_model_ = metadata["onnx_model_"]
         obj.output_quantizers = metadata["output_quantizers"]
+        obj._fhe_ensembling = metadata["_fhe_ensembling"]
         obj._tree_inference = tree_to_numpy(
             obj.sklearn_model,
             numpy.zeros((len(obj.input_quantizers),))[None, ...],
             framework=obj.framework,
-            output_n_bits=obj.n_bits,
+            output_n_bits=obj.n_bits["op_leaves"] if isinstance(obj.n_bits, Dict) else obj.n_bits,
+            fhe_ensembling=obj._fhe_ensembling,
         )[0]
         obj.post_processing_params = metadata["post_processing_params"]
 
@@ -162,7 +165,7 @@ def __init__(
         max_leaf_nodes=None,
         min_impurity_decrease=0.0,
         ccp_alpha=0.0,
-        n_bits: int = 6,
+        n_bits: Union[int, Dict[str, int]] = 6,
     ):
         """Initialize the DecisionTreeRegressor.
 
@@ -208,6 +211,7 @@ def dump_dict(self) -> Dict[str, Any]:
         metadata["onnx_model_"] = self.onnx_model_
         metadata["framework"] = self.framework
         metadata["post_processing_params"] = self.post_processing_params
+        metadata["_fhe_ensembling"] = self._fhe_ensembling
 
         # Scikit-Learn
         metadata["criterion"] = self.criterion
@@ -233,16 +237,19 @@ def load_dict(cls, metadata: Dict):
         # Concrete-ML
         obj.sklearn_model = metadata["sklearn_model"]
         obj._is_fitted = metadata["_is_fitted"]
+        obj._fhe_ensembling = metadata["_fhe_ensembling"]
         obj._is_compiled = metadata["_is_compiled"]
         obj.input_quantizers = metadata["input_quantizers"]
         obj.framework = metadata["framework"]
         obj.onnx_model_ = metadata["onnx_model_"]
         obj.output_quantizers = metadata["output_quantizers"]
+        obj._fhe_ensembling = metadata["_fhe_ensembling"]
         obj._tree_inference = tree_to_numpy(
             obj.sklearn_model,
             numpy.zeros((len(obj.input_quantizers),))[None, ...],
             framework=obj.framework,
-            output_n_bits=obj.n_bits,
+            output_n_bits=obj.n_bits["op_leaves"] if isinstance(obj.n_bits, Dict) else obj.n_bits,
+            fhe_ensembling=obj._fhe_ensembling,
         )[0]
         obj.post_processing_params = metadata["post_processing_params"]