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
This commit is contained in:
youben11
2021-12-09 16:42:41 +01:00
committed by Ayoub Benaissa
parent d34792c977
commit 550318f67e
2 changed files with 27 additions and 3 deletions

View File

@@ -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)

View File

@@ -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>