From 550318f67e59a2e2700b2a67abfbec04ccfcb50f Mon Sep 17 00:00:00 2001
From: youben11 <ayoub.benaissa@zama.ai>
Date: Thu, 9 Dec 2021 16:42:41 +0100
Subject: [PATCH] feat: support different numpy object types as input

- Scalar types: int, np.uint8 (to extend with other types later), and np.ndarray with shape == ()
- Tensor types: np.ndarray
---
 .../lib/Bindings/Python/zamalang/compiler.py  |  8 ++++---
 compiler/tests/python/test_compiler_engine.py | 22 +++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/compiler/lib/Bindings/Python/zamalang/compiler.py b/compiler/lib/Bindings/Python/zamalang/compiler.py
index 927ee5cef..87e8fbaa9 100644
--- a/compiler/lib/Bindings/Python/zamalang/compiler.py
+++ b/compiler/lib/Bindings/Python/zamalang/compiler.py
@@ -95,9 +95,9 @@ def create_execution_argument(value: Union[int, np.ndarray]) -> "_LambdaArgument
     Returns:
         _LambdaArgument: lambda argument holding the appropriate value
     """
-    if not isinstance(value, (int, np.ndarray)):
-        raise TypeError("value of execution argument must be either int or numpy.array")
-    if isinstance(value, int):
+    if not isinstance(value, (int, np.ndarray, np.uint8)):
+        raise TypeError("value of execution argument must be either int, numpy.array or numpy.uint8")
+    if isinstance(value, (int, np.uint8)):
         if not (0 <= value < (2 ** 64 - 1)):
             raise TypeError(
                 "single integer must be in the range [0, 2**64 - 1] (uint64)"
@@ -105,6 +105,8 @@ def create_execution_argument(value: Union[int, np.ndarray]) -> "_LambdaArgument
         return _LambdaArgument.from_scalar(value)
     else:
         assert isinstance(value, np.ndarray)
+        if value.shape == ():
+            return _LambdaArgument.from_scalar(value)
         if value.dtype != np.uint8:
             raise TypeError("numpy.array must be of dtype uint8")
         return _LambdaArgument.from_tensor(value.flatten().tolist(), value.shape)
diff --git a/compiler/tests/python/test_compiler_engine.py b/compiler/tests/python/test_compiler_engine.py
index 7ad500d20..1eca9b4be 100644
--- a/compiler/tests/python/test_compiler_engine.py
+++ b/compiler/tests/python/test_compiler_engine.py
@@ -21,6 +21,28 @@ KEY_SET_CACHE_PATH =  os.path.join(tempfile.gettempdir(), 'KeySetCache')
             12,
             id="add_eint_int",
         ),
+        pytest.param(
+            """
+            func @main(%arg0: !HLFHE.eint<7>, %arg1: i8) -> !HLFHE.eint<7> {
+                %1 = "HLFHE.add_eint_int"(%arg0, %arg1): (!HLFHE.eint<7>, i8) -> (!HLFHE.eint<7>)
+                return %1: !HLFHE.eint<7>
+            }
+            """,
+            (np.array(4, dtype=np.uint8), np.array(5, dtype=np.uint8)),
+            9,
+            id="add_eint_int_with_ndarray_as_scalar",
+        ),
+        pytest.param(
+            """
+            func @main(%arg0: !HLFHE.eint<7>, %arg1: i8) -> !HLFHE.eint<7> {
+                %1 = "HLFHE.add_eint_int"(%arg0, %arg1): (!HLFHE.eint<7>, i8) -> (!HLFHE.eint<7>)
+                return %1: !HLFHE.eint<7>
+            }
+            """,
+            (np.uint8(3), np.uint8(5)),
+            8,
+            id="add_eint_int_with_np_uint8_as_scalar",
+        ),
         pytest.param(
             """
             func @main(%arg0: tensor<4x!HLFHE.eint<7>>, %arg1: tensor<4xi8>) -> !HLFHE.eint<7>