chore: add check proper keyword for rounding_threshold_bits (#624)

src/concrete/ml/common/utils.py

@@ -623,6 +623,7 @@ def process_rounding_threshold_bits(rounding_threshold_bits):
     Raises:
         NotImplementedError: If 'auto' rounding is specified but not implemented.
         ValueError: If an invalid type or value is provided for rounding_threshold_bits.
+        KeyError: If the dict contains keys other than 'n_bits' and 'method'.
     """
     n_bits_rounding: Union[None, str, int] = None
     method: Exactness = Exactness.EXACT
@@ -632,9 +633,16 @@ def process_rounding_threshold_bits(rounding_threshold_bits):
         if isinstance(rounding_threshold_bits, int):
             n_bits_rounding = rounding_threshold_bits
         elif isinstance(rounding_threshold_bits, dict):
+            valid_keys = {"n_bits", "method"}
+            if not valid_keys.issuperset(rounding_threshold_bits.keys()):
+                raise KeyError(
+                    f"Invalid keys in rounding_threshold_bits. Allowed keys are {valid_keys}."
+                )
             n_bits_rounding = rounding_threshold_bits.get("n_bits")
             if n_bits_rounding == "auto":
                 raise NotImplementedError("Automatic rounding is not implemented yet.")
+            if not isinstance(n_bits_rounding, int):
+                raise ValueError("n_bits must be an integer.")
             method = rounding_threshold_bits.get("method", method)
             if not isinstance(method, Exactness):
                 method_str = method.upper()

tests/torch/test_compile_torch.py

@@ -1406,7 +1406,7 @@ def forward(self, x):
             "Invalid type for rounding_threshold_bits. Must be int or dict.",
         ),
         (
-            {"method": "INVALID_METHOD"},
+            {"n_bits": 4, "method": "INVALID_METHOD"},
             ValueError,
             "INVALID_METHOD is not a valid method. Must be one of EXACT, APPROXIMATE.",
         ),
@@ -1420,6 +1420,16 @@ def forward(self, x):
             ValueError,
             "n_bits_rounding must be between 2 and 8 inclusive",
         ),
+        (
+            {"invalid_key": 4},
+            KeyError,
+            "Invalid keys in rounding_threshold_bits. Allowed keys are {'n_bits', 'method'}.",
+        ),
+        (
+            {"n_bits": "not_an_int"},
+            ValueError,
+            "n_bits must be an integer.",
+        ),
     ],
 )
 def test_compile_torch_model_rounding_threshold_bits_errors(