From 1c4ebc06d82f9f8a40562049322f758717e60b52 Mon Sep 17 00:00:00 2001 From: kcelia Date: Mon, 29 Jan 2024 16:30:02 +0100 Subject: [PATCH] chore: update --- src/concrete/ml/quantization/post_training.py | 2 ++ tests/sklearn/test_sklearn_models.py | 16 ++++++++-------- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/concrete/ml/quantization/post_training.py b/src/concrete/ml/quantization/post_training.py index 46bed0214..022508507 100644 --- a/src/concrete/ml/quantization/post_training.py +++ b/src/concrete/ml/quantization/post_training.py @@ -86,6 +86,8 @@ def _inspect_tree_n_bits(n_bits): ) +# Find a better naming to describe leaf quantization in tree-based models +# FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/4258 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. diff --git a/tests/sklearn/test_sklearn_models.py b/tests/sklearn/test_sklearn_models.py index 58423df78..6666087e1 100644 --- a/tests/sklearn/test_sklearn_models.py +++ b/tests/sklearn/test_sklearn_models.py @@ -1210,7 +1210,11 @@ def check_fhe_sum_for_tree_based_models( # Run the test with more samples during weekly CIs if is_weekly_option: - fhe_test = get_random_samples(x, n_sample=5) + fhe_samples = 5 + else: + fhe_samples = 1 + + fhe_test = get_random_samples(x, n_sample=fhe_samples) # pylint: disable=protected-access assert not model._fhe_ensembling, "`_fhe_ensembling` is disabled by default." @@ -1218,14 +1222,11 @@ def check_fhe_sum_for_tree_based_models( non_fhe_sum_predict_quantized = predict_method(x, fhe="disable") non_fhe_sum_predict_simulate = predict_method(x, fhe="simulate") + non_fhe_sum_predict_fhe = predict_method(fhe_test, fhe="execute") # Sanity check array_allclose_and_same_shape(non_fhe_sum_predict_quantized, non_fhe_sum_predict_simulate) - # Compute the FHE predictions only during weekly CIs - if is_weekly_option: - non_fhe_sum_predict_fhe = predict_method(fhe_test, fhe="execute") - # pylint: disable=protected-access model._fhe_ensembling = True @@ -1233,6 +1234,7 @@ def check_fhe_sum_for_tree_based_models( fhe_sum_predict_quantized = predict_method(x, fhe="disable") fhe_sum_predict_simulate = predict_method(x, fhe="simulate") + fhe_sum_predict_fhe = predict_method(fhe_test, fhe="execute") # Sanity check array_allclose_and_same_shape(fhe_sum_predict_quantized, fhe_sum_predict_simulate) @@ -1240,9 +1242,7 @@ def check_fhe_sum_for_tree_based_models( # Check that we have the exact same predictions array_allclose_and_same_shape(fhe_sum_predict_quantized, non_fhe_sum_predict_quantized) array_allclose_and_same_shape(fhe_sum_predict_simulate, non_fhe_sum_predict_simulate) - if is_weekly_option: - fhe_sum_predict_fhe = predict_method(fhe_test, fhe="execute") - array_allclose_and_same_shape(fhe_sum_predict_fhe, non_fhe_sum_predict_fhe) + array_allclose_and_same_shape(fhe_sum_predict_fhe, non_fhe_sum_predict_fhe) # Neural network models are skipped for this test