From 39b168dfd594f7648073a72ccb2d71827ed8cbaa Mon Sep 17 00:00:00 2001
From: Benoit Chevallier-Mames <benoit.chevalliermames@zama.ai>
Date: Wed, 24 Nov 2021 15:28:34 +0100
Subject: [PATCH] fix: fix issues for certain form of functions

Table which was created contained too large values, for cells which were never accessed
closes #949
---
 concrete/common/mlir/node_converter.py | 7 ++++++-
 tests/numpy/test_compile.py            | 3 +--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/concrete/common/mlir/node_converter.py b/concrete/common/mlir/node_converter.py
index d8457b6c6..78b4edcb7 100644
--- a/concrete/common/mlir/node_converter.py
+++ b/concrete/common/mlir/node_converter.py
@@ -282,7 +282,12 @@ class IntermediateNodeConverter:
             table = tables[0][0]
 
             lut_shape: Tuple[int, ...] = (len(table),)
-            lut_values = numpy.array(table, dtype=numpy.uint64)
+
+            # The reduction on 63b is to avoid problems like doing a TLU of
+            # the form T[j] = 2<<j, for j which is supposed to be 7b as per
+            # constraint of the compiler, while in practice, it is a small
+            # value. Reducing on 64b was not ok for some reason
+            lut_values = numpy.array(table % (2 << 63), dtype=numpy.uint64)
         else:
             assert_true(isinstance(output, TensorValue))
             assert isinstance(output, TensorValue)
diff --git a/tests/numpy/test_compile.py b/tests/numpy/test_compile.py
index b8c5dd355..ae63a463f 100644
--- a/tests/numpy/test_compile.py
+++ b/tests/numpy/test_compile.py
@@ -683,9 +683,8 @@ def test_compile_and_run_correctness(
     [
         pytest.param(lambda x: x ** 2, ((0, 10),), ["x"]),
         pytest.param(lambda x: 2 ** (x % 5), ((0, 20),), ["x"]),
-        # FIXME: Fails, #949 pytest.param(lambda x: abs(~x), ((0, 13),), ["x"]),
         pytest.param(lambda x: x << 1, ((0, 13),), ["x"]),
-        # FIXME: Fails, #949 pytest.param(lambda x: 2 << (x % 6), ((0, 13),), ["x"]),
+        pytest.param(lambda x: 2 << (x % 6), ((0, 13),), ["x"]),
         pytest.param(lambda x: x >> 2, ((30, 100),), ["x"]),
         pytest.param(lambda x: 115 >> (x % 3), ((0, 17),), ["x"]),
         pytest.param(lambda x: x % 7, ((0, 100),), ["x"]),