fix(frontend-python): bits with direct circuit, single bit is precisi…

…on 1
zama-ai · Apr 3, 2024 · aa3b4fa · aa3b4fa
1 parent 15b2291
commit aa3b4fa
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/frontends/concrete-python/concrete/fhe/extensions/bits.py b/frontends/concrete-python/concrete/fhe/extensions/bits.py
@@ -132,10 +132,18 @@ def evaluator(x, bits):  # pylint: disable=redefined-outer-name
             return result
 
         if isinstance(self.value, Tracer):
+            output_value = deepcopy(self.value.output)
+            direct_single_bit = (
+                Tracer._is_direct
+                and isinstance(index, int)
+                and isinstance(output_value.dtype, Integer)
+            )
+            if direct_single_bit:
+                output_value.dtype.bit_width = 1  # type: ignore[attr-defined]
             computation = Node.generic(
                 "extract_bit_pattern",
                 [deepcopy(self.value.output)],
-                deepcopy(self.value.output),
+                output_value,
                 evaluator,
                 kwargs={"bits": index},
             )