mplang-nightly 0.1.dev163__tar.gz → 0.1.dev165__tar.gz

{mplang_nightly-0.1.dev163 → mplang_nightly-0.1.dev165}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev163
+Version: 0.1.dev165
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License

{mplang_nightly-0.1.dev163 → mplang_nightly-0.1.dev165}/mplang/core/expr/evaluator.py

@@ -27,6 +27,8 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Any, Protocol
 
+import numpy as np
+
 from mplang.core.comm import ICommunicator
 from mplang.core.expr.ast import (
     AccessExpr,
@@ -47,6 +49,7 @@ from mplang.core.expr.walk import walk_dataflow
 from mplang.core.mask import Mask
 from mplang.core.pfunc import PFunction
 from mplang.kernels.context import RuntimeContext
+from mplang.kernels.value import Value
 
 
 class IEvaluator(Protocol):
@@ -149,12 +152,12 @@ class EvalSemantic:
     def _as_optional_int(val: Any) -> int | None:
         """Convert a value to int if possible, preserving None.
 
-        Handles Python ints, numpy scalars with .item(), and None.
+        Handles Python ints, floats, numpy scalar types (e.g., np.int32, np.float64), and None.
+        Uses int(val) for conversion which works with numpy scalars via __int__().
         """
+        val = EvalSemantic._unwrap_value(val)
         if val is None:
             return None
-        if hasattr(val, "item"):
-            return int(val.item())
         return int(val)
 
     def _simple_allgather(self, value: Any) -> list[Any]:
@@ -167,6 +170,7 @@ class EvalSemantic:
         Returns a list of length world_size with entries ordered by rank.
         """
         ws = self.comm.world_size
+        value = self._unwrap_value(value)
         # Trivial fast-path
         if ws == 1:
             return [value]
@@ -185,7 +189,12 @@ class EvalSemantic:
 
     def _verify_uniform_predicate(self, pred: Any) -> None:
         # Runtime uniformity check (O(P^2) send/recv emulation).
-        vals = self._simple_allgather(bool(pred))
+        # Use Value.to_bool() if available, otherwise unwrap and convert
+        if isinstance(pred, Value):
+            pred_bool = pred.to_bool()
+        else:
+            pred_bool = bool(self._unwrap_value(pred))
+        vals = self._simple_allgather(pred_bool)
         if not vals:
             raise ValueError("uniform_cond: empty gather for predicate")
         first = vals[0]
@@ -209,13 +218,33 @@ class EvalSemantic:
         assert len(cond_result) == 1, (
             f"Condition function must return a single value, got {cond_result}"
         )
-        cond_value = cond_result[0]
-        if cond_value is None:
+        cond_val = cond_result[0]
+        if cond_val is None:
             raise RuntimeError(
                 "while_loop condition produced None on rank "
                 f"{self.rank}; ensure the predicate yields a boolean for every party."
             )
-        return cond_value
+        # Use Value.to_bool() if available for cleaner conversion
+        if isinstance(cond_val, Value):
+            return cond_val.to_bool()
+        return bool(self._unwrap_value(cond_val))
+
+    @staticmethod
+    def _unwrap_value(value: Any) -> Any:
+        """Convert Value payloads to numpy/python equivalents when possible."""
+        if value is None:
+            return None
+        if isinstance(value, Value):
+            # Try to_numpy first for broader compatibility
+            to_numpy = getattr(value, "to_numpy", None)
+            if callable(to_numpy):
+                arr = to_numpy()
+                if isinstance(arr, np.ndarray):
+                    if arr.size == 1:
+                        return arr.item()
+                    return arr
+                return arr
+        return value
 
 
 class RecursiveEvaluator(EvalSemantic, ExprVisitor):
@@ -296,15 +325,21 @@ class RecursiveEvaluator(EvalSemantic, ExprVisitor):
           * Add optional static uniform inference (data provenance) to elide the
             runtime check when predicate uniformity is provable at trace time.
         """
-        pred = self._value(expr.pred)
-        if pred is None:
+        pred_val = self._value(expr.pred)
+        if pred_val is None:
             return [None] * len(expr.mptypes)
 
         if expr.verify_uniform:
-            self._verify_uniform_predicate(pred)
+            self._verify_uniform_predicate(pred_val)
+
+        # Convert to bool using Value.to_bool() if available
+        if isinstance(pred_val, Value):
+            pred = pred_val.to_bool()
+        else:
+            pred = bool(self._unwrap_value(pred_val))
 
         # Only evaluate selected branch locally
-        if pred:
+        if bool(pred):
             then_call = CallExpr(expr.then_fn, expr.args)
             return self._values(then_call)
         else:
@@ -435,15 +470,20 @@ class IterativeEvaluator(EvalSemantic):
                 res = self._iter_eval_graph(node.fn.body, {**env, **sub_env})
                 symbols[id(node)] = res
             elif isinstance(node, CondExpr):
-                pred_v = self._first(symbols[id(node.pred)])
+                pred_val = self._first(symbols[id(node.pred)])
                 arg_vals = [self._first(symbols[id(a)]) for a in node.args]
-                if pred_v is None:
+                if pred_val is None:
                     symbols[id(node)] = [None] * len(node.mptypes)
                 else:
                     # Optional uniform verification identical to recursive evaluator (DRY helper).
                     if node.verify_uniform:
-                        self._verify_uniform_predicate(pred_v)
-                    if bool(pred_v):
+                        self._verify_uniform_predicate(pred_val)
+                    # Convert to bool using Value.to_bool() if available
+                    if isinstance(pred_val, Value):
+                        pred = pred_val.to_bool()
+                    else:
+                        pred = bool(self._unwrap_value(pred_val))
+                    if pred:
                         sub_env = dict(zip(node.then_fn.params, arg_vals, strict=True))
                         res = self._iter_eval_graph(
                             node.then_fn.body, {**env, **sub_env}

{mplang_nightly-0.1.dev163 → mplang_nightly-0.1.dev165}/mplang/device.py

@@ -207,15 +207,8 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
         assert len(frm_dev.members) == 1 and len(to_dev.members) == 1
         frm_rank = frm_dev.members[0].rank
         tee_rank = to_dev.members[0].rank
-        platform = to_dev.config.get("platform")
-        if not platform:
-            raise ValueError(
-                f"TEE device '{to_dev_id}' is missing 'platform' in its config."
-            )
         # Ensure sessions (both directions) exist for this PPU<->TEE pair
-        sess_p, sess_t = _ensure_tee_session(
-            frm_dev_id, to_dev_id, frm_rank, tee_rank, platform
-        )
+        sess_p, sess_t = _ensure_tee_session(frm_dev_id, to_dev_id, frm_rank, tee_rank)
         # Bytes-only path: pack -> enc -> p2p -> dec -> unpack (with static out type)
         obj_ty = TensorType.from_obj(obj)
         b = simp.runAt(frm_rank, builtin.pack)(obj)
@@ -229,15 +222,8 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
         assert len(frm_dev.members) == 1 and len(to_dev.members) == 1
         tee_rank = frm_dev.members[0].rank
         ppu_rank = to_dev.members[0].rank
-        platform = frm_dev.config.get("platform")
-        if not platform:
-            raise ValueError(
-                f"TEE device '{frm_dev_id}' is missing 'platform' in its config."
-            )
         # Ensure bidirectional session established for this pair
-        sess_p, sess_t = _ensure_tee_session(
-            to_dev_id, frm_dev_id, ppu_rank, tee_rank, platform
-        )
+        sess_p, sess_t = _ensure_tee_session(to_dev_id, frm_dev_id, ppu_rank, tee_rank)
         obj_ty = TensorType.from_obj(obj)
         b = simp.runAt(tee_rank, builtin.pack)(obj)
         ct = simp.runAt(tee_rank, crypto.enc)(b, sess_t)
@@ -259,7 +245,7 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
 
 
 def _ensure_tee_session(
-    frm_dev_id: str, to_dev_id: str, frm_rank: int, tee_rank: int, platform: str
+    frm_dev_id: str, to_dev_id: str, frm_rank: int, tee_rank: int
 ) -> tuple[MPObject, MPObject]:
     """Ensure a TEE session (sess_p at sender, sess_t at TEE) exists.
 
@@ -281,7 +267,7 @@ def _ensure_tee_session(
 
     # 2) Send quote to sender and attest to obtain TEE pk
     quote_at_sender = mpi.p2p(tee_rank, frm_rank, quote)
-    tee_pk_at_sender = simp.runAt(frm_rank, tee.attest)(quote_at_sender, platform)
+    tee_pk_at_sender = simp.runAt(frm_rank, tee.attest)(quote_at_sender)
 
     # 3) Sender generates its ephemeral keypair and sends its pk to TEE
     v_sk, v_pk = simp.runAt(frm_rank, crypto.kem_keygen)(_TEE_KEM_SUITE)

mplang_nightly-0.1.dev165/mplang/kernels/__init__.py

@@ -0,0 +1,41 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from mplang.kernels.value import (
+    BytesBlob,
+    TableValue,
+    TensorValue,
+    Value,
+    ValueDecodeError,
+    ValueError,
+    decode_value,
+    encode_value,
+    is_value_envelope,
+    list_value_kinds,
+    register_value,
+)
+
+__all__ = [
+    "BytesBlob",
+    "TableValue",
+    "TensorValue",
+    "Value",
+    "ValueDecodeError",
+    "ValueError",
+    "decode_value",
+    "encode_value",
+    "is_value_envelope",
+    "list_value_kinds",
+    "register_value",
+]

{mplang_nightly-0.1.dev163 → mplang_nightly-0.1.dev165}/mplang/kernels/builtin.py

@@ -14,38 +14,25 @@
 
 from __future__ import annotations
 
-from typing import Any
-
 import numpy as np
-import pandas as pd
 
 from mplang.core.pfunc import PFunction
 from mplang.core.table import TableType
 from mplang.core.tensor import TensorType
 from mplang.kernels.base import cur_kctx, kernel_def
+from mplang.kernels.value import TableValue, TensorValue, Value
 from mplang.runtime.data_providers import get_provider, resolve_uri
 from mplang.utils import table_utils
 
 
-def _to_numpy(obj: Any) -> np.ndarray:  # minimal helper to avoid duplicating logic
-    if isinstance(obj, np.ndarray):
-        return obj
-    if hasattr(obj, "numpy"):
-        try:
-            return np.asarray(obj.numpy())  # type: ignore
-        except Exception:
-            pass
-    return np.asarray(obj)
-
-
 @kernel_def("builtin.identity")
-def _identity(pfunc: PFunction, value: Any) -> Any:
+def _identity(pfunc: PFunction, value: Value) -> Value:
     # Runtime guarantees exactly one argument; no extra arity checks here.
     return value
 
 
 @kernel_def("builtin.read")
-def _read(pfunc: PFunction) -> Any:
+def _read(pfunc: PFunction) -> Value:
     path = pfunc.attrs.get("path")
     if path is None:
         raise ValueError("missing path attr for builtin.read")
@@ -56,13 +43,25 @@ def _read(pfunc: PFunction) -> Any:
         raise NotImplementedError(f"no resource provider for scheme: {uri.scheme}")
     ctx = cur_kctx()
     try:
-        return prov.read(uri, out_t, ctx=ctx)
+        data = prov.read(uri, out_t, ctx=ctx)
     except Exception as e:  # pragma: no cover - provider errors
         raise RuntimeError(f"builtin.read failed: {e}") from e
 
+    if isinstance(out_t, TableType):
+        if isinstance(data, TableValue):
+            return data
+        return TableValue(data)
+    if isinstance(out_t, TensorType):
+        if isinstance(data, TensorValue):
+            return data
+        return TensorValue(np.asarray(data))
+    raise TypeError(
+        f"builtin.read only supports TableType/TensorType outputs, got {type(out_t).__name__}"
+    )
+
 
 @kernel_def("builtin.write")
-def _write(pfunc: PFunction, obj: Any) -> Any:
+def _write(pfunc: PFunction, obj: Value) -> Value:
     path = pfunc.attrs.get("path")
     if path is None:
         raise ValueError("missing path attr for builtin.write")
@@ -70,16 +69,18 @@ def _write(pfunc: PFunction, obj: Any) -> Any:
     prov = get_provider(uri.scheme)
     if prov is None:
         raise NotImplementedError(f"no resource provider for scheme: {uri.scheme}")
+    # Pass Value object directly to provider - let provider decide how to handle it
     ctx = cur_kctx()
     try:
         prov.write(uri, obj, ctx=ctx)
-        return obj
     except Exception as e:  # pragma: no cover
         raise RuntimeError(f"builtin.write failed: {e}") from e
+    return obj
 
 
 @kernel_def("builtin.constant")
-def _constant(pfunc: PFunction) -> Any:
+def _constant(pfunc: PFunction) -> Value:
+    """Return constants as Value types (TensorValue or TableValue)."""
     data_bytes = pfunc.attrs.get("data_bytes")
     if data_bytes is None:
         raise ValueError("missing data_bytes attr for builtin.constant")
@@ -89,69 +90,86 @@ def _constant(pfunc: PFunction) -> Any:
         if fmt != "bytes[csv]":
             raise ValueError(f"unsupported table constant format {fmt}")
         df = table_utils.csv_to_dataframe(data_bytes)
-        return df
+        return TableValue(df)
     # tensor path
     shape = out_t.shape  # type: ignore[attr-defined,union-attr]
     dtype = out_t.dtype.numpy_dtype()  # type: ignore[attr-defined,union-attr]
     arr = np.frombuffer(data_bytes, dtype=dtype).reshape(shape)
-    return arr
+    return TensorValue(arr)
 
 
 @kernel_def("builtin.rank")
-def _rank(pfunc: PFunction) -> Any:
+def _rank(pfunc: PFunction) -> TensorValue:
+    """Return rank as TensorValue."""
     ctx = cur_kctx()
-    return np.array(ctx.rank, dtype=np.uint64)
+    arr = np.array(ctx.rank, dtype=np.uint64)
+    return TensorValue(arr)
 
 
 @kernel_def("builtin.prand")
-def _prand(pfunc: PFunction) -> Any:
+def _prand(pfunc: PFunction) -> TensorValue:
+    """Return random data as TensorValue."""
     shape = pfunc.attrs.get("shape", ())
     rng = np.random.default_rng()
     info = np.iinfo(np.uint64)
     data = rng.integers(
         low=info.min, high=info.max, size=shape, dtype=np.uint64, endpoint=True
     )
-    return data
+    return TensorValue(data)
 
 
 @kernel_def("builtin.table_to_tensor")
-def _table_to_tensor(pfunc: PFunction, table: Any) -> Any:
-    if not isinstance(table, pd.DataFrame):
-        raise TypeError("expected pandas DataFrame")
-    if table.shape[1] == 0:
+def _table_to_tensor(pfunc: PFunction, table: TableValue) -> TensorValue:
+    """Convert table to tensor, return as TensorValue."""
+    arrow_table = table.to_arrow()
+    if arrow_table.num_columns == 0:
         raise ValueError("cannot pack empty table")
-    mat = np.column_stack([table[col].to_numpy() for col in table.columns])
-    return mat
+    # Convert Arrow columns to numpy arrays and stack
+    mat = np.column_stack([
+        arrow_table.column(i).to_numpy() for i in range(arrow_table.num_columns)
+    ])
+    return TensorValue(mat)
 
 
 @kernel_def("builtin.tensor_to_table")
-def _tensor_to_table(pfunc: PFunction, tensor: Any) -> Any:
-    arr = _to_numpy(tensor)
+def _tensor_to_table(pfunc: PFunction, tensor: TensorValue) -> TableValue:
+    """Convert tensor to table, return as TableValue."""
+    import pyarrow as pa  # type: ignore
+
+    arr = tensor.to_numpy()
     if arr.ndim != 2:
         raise ValueError("tensor_to_table expects rank-2 array")
     col_names = pfunc.attrs.get("column_names")
     if col_names is None:
         raise ValueError("missing column_names attr")
-    df = pd.DataFrame(arr, columns=list(col_names))
-    return df
+    # Create Arrow table directly from numpy array columns
+    arrays = [pa.array(arr[:, i]) for i in range(arr.shape[1])]
+    arrow_table = pa.table(dict(zip(col_names, arrays, strict=True)))
+    return TableValue(arrow_table)
 
 
-def _summ(v: Any) -> str:
+def _summ(v: Value) -> str:
     try:
-        if isinstance(v, pd.DataFrame):
-            return str(v.head(8).to_string(index=False))
-        arr = _to_numpy(v)
-        return str(
-            np.array2string(
-                arr, threshold=64, edgeitems=3, precision=6, suppress_small=True
+        if isinstance(v, TableValue):
+            # Use Arrow's native string representation (more efficient)
+            arrow_table = v.to_arrow()
+            # Show first 8 rows
+            preview = arrow_table.slice(0, min(8, arrow_table.num_rows))
+            return str(preview)
+        if isinstance(v, TensorValue):
+            arr = v.to_numpy()
+            return str(
+                np.array2string(
+                    arr, threshold=64, edgeitems=3, precision=6, suppress_small=True
+                )
             )
-        )
+        return repr(v)
     except Exception as e:  # pragma: no cover
         return f"<unprintable {type(v).__name__}: {e}>"
 
 
 @kernel_def("builtin.debug_print")
-def _debug_print(pfunc: PFunction, val: Any) -> Any:
+def _debug_print(pfunc: PFunction, val: Value) -> Value:
     prefix = pfunc.attrs.get("prefix", "")
     ctx = cur_kctx()
     print(f"[debug_print][rank={ctx.rank}] {prefix}{_summ(val)}")
@@ -159,7 +177,7 @@ def _debug_print(pfunc: PFunction, val: Any) -> Any:
 
 
 @kernel_def("builtin.pack")
-def _pack(pfunc: PFunction, value: Any) -> Any:
+def _pack(pfunc: PFunction, value: Value) -> TensorValue:
     outs_info = pfunc.outs_info
     if len(outs_info) != 1:
         raise ValueError("builtin.pack expects single output type")
@@ -169,22 +187,33 @@ def _pack(pfunc: PFunction, value: Any) -> Any:
     if out_ty.dtype.numpy_dtype() != np.uint8:
         raise TypeError("builtin.pack output dtype must be uint8")
 
-    if isinstance(value, pd.DataFrame):
-        csv_bytes = table_utils.dataframe_to_csv(value)
-        return np.frombuffer(csv_bytes, dtype=np.uint8)
+    if isinstance(value, TableValue):
+        # Serialize Arrow table using IPC stream for consistency with Value serde
+        import pyarrow as pa  # type: ignore
+        import pyarrow.ipc as pa_ipc  # type: ignore
+
+        arrow_table = value.to_arrow()
+        sink = pa.BufferOutputStream()
+        with pa_ipc.new_stream(sink, arrow_table.schema) as writer:  # type: ignore[arg-type]
+            writer.write_table(arrow_table)  # type: ignore[arg-type]
+        ipc_bytes = sink.getvalue().to_pybytes()
+        return TensorValue(np.frombuffer(ipc_bytes, dtype=np.uint8))
+
+    if isinstance(value, TensorValue):
+        arr = value.to_numpy()
+        return TensorValue(np.frombuffer(arr.tobytes(order="C"), dtype=np.uint8))
 
-    arr = _to_numpy(value)
-    return np.frombuffer(arr.tobytes(order="C"), dtype=np.uint8)
+    raise TypeError(f"builtin.pack does not support Value type {type(value).__name__}")
 
 
 @kernel_def("builtin.unpack")
-def _unpack(pfunc: PFunction, packed: Any) -> Any:
+def _unpack(pfunc: PFunction, packed: TensorValue) -> Value:
     outs_info = pfunc.outs_info
     if len(outs_info) != 1:
         raise ValueError("builtin.unpack expects single output type")
     out_ty = outs_info[0]
 
-    b = np.asarray(packed, dtype=np.uint8).reshape(-1)
+    b = packed.to_numpy().astype(np.uint8, copy=False).reshape(-1)
 
     if isinstance(out_ty, TensorType):
         np_dtype = out_ty.dtype.numpy_dtype()
@@ -198,10 +227,16 @@ def _unpack(pfunc: PFunction, packed: Any) -> Any:
                 f"unpack size mismatch: got {b.size} bytes, expect {expected} for {np_dtype} {shape}"
             )
         arr = np.frombuffer(b.tobytes(), dtype=np_dtype)
-        return arr.reshape(shape)
+        return TensorValue(arr.reshape(shape))
 
     if isinstance(out_ty, TableType):
-        csv_bytes = b.tobytes()
-        return table_utils.csv_to_dataframe(csv_bytes)
+        # Deserialize Arrow IPC stream back to TableValue
+        import pyarrow as pa  # type: ignore
+        import pyarrow.ipc as pa_ipc  # type: ignore
+
+        buf = pa.py_buffer(b.tobytes())
+        reader = pa_ipc.open_stream(buf)
+        table = reader.read_all()
+        return TableValue(table)
 
     raise TypeError("builtin.unpack output type must be TensorType or TableType")

{mplang_nightly-0.1.dev163 → mplang_nightly-0.1.dev165}/mplang/kernels/crypto.py

@@ -15,15 +15,15 @@
 from __future__ import annotations
 
 import os
-from typing import Any
 
 import numpy as np
 
 from mplang.core.pfunc import PFunction
 from mplang.kernels.base import cur_kctx, kernel_def
+from mplang.kernels.value import TensorValue
 from mplang.utils.crypto import blake2b
 
-__all__: list[str] = []  # flat kernels only
+__all__: list[str] = []  # No public exports currently
 
 
 def _get_rng() -> np.random.Generator:
@@ -54,62 +54,71 @@ def _keystream(key: bytes, nonce: bytes, length: int) -> bytes:
 
 
 @kernel_def("crypto.keygen")
-def _crypto_keygen(pfunc: PFunction) -> Any:
+def _crypto_keygen(pfunc: PFunction) -> TensorValue:
     length = int(pfunc.attrs.get("length", 32))
     rng = _get_rng()
     key = rng.integers(0, 256, size=(length,), dtype=np.uint8)
-    return key
+    return TensorValue(key)
 
 
 @kernel_def("crypto.enc")
-def _crypto_encrypt(pfunc: PFunction, pt_bytes: Any, key: Any) -> Any:
-    pt_bytes = np.asarray(pt_bytes, dtype=np.uint8)
-    key = np.asarray(key, dtype=np.uint8)
+def _crypto_encrypt(
+    pfunc: PFunction, pt_bytes: TensorValue, key: TensorValue
+) -> TensorValue:
+    pt_bytes_np = pt_bytes.to_numpy().astype(np.uint8, copy=False)
+    key_np = key.to_numpy().astype(np.uint8, copy=False)
     rng = _get_rng()
     nonce = rng.integers(0, 256, size=(12,), dtype=np.uint8)
     stream = np.frombuffer(
-        _keystream(key.tobytes(), nonce.tobytes(), pt_bytes.size), dtype=np.uint8
+        _keystream(key_np.tobytes(), nonce.tobytes(), pt_bytes_np.size), dtype=np.uint8
     )
-    ct = (pt_bytes ^ stream).astype(np.uint8)
+    ct = (pt_bytes_np ^ stream).astype(np.uint8)
     out = np.concatenate([nonce, ct]).astype(np.uint8)
-    return out
+    return TensorValue(out)
 
 
 @kernel_def("crypto.dec")
-def _crypto_decrypt(pfunc: PFunction, ct_with_nonce: Any, key: Any) -> Any:
-    ct_with_nonce = np.asarray(ct_with_nonce, dtype=np.uint8)
-    key = np.asarray(key, dtype=np.uint8)
-    nonce = ct_with_nonce[:12]
-    ct = ct_with_nonce[12:]
+def _crypto_decrypt(
+    pfunc: PFunction, ct_with_nonce: TensorValue, key: TensorValue
+) -> TensorValue:
+    ct_np = ct_with_nonce.to_numpy().astype(np.uint8, copy=False)
+    key_np = key.to_numpy().astype(np.uint8, copy=False)
+    nonce = ct_np[:12]
+    ct = ct_np[12:]
     stream = np.frombuffer(
-        _keystream(key.tobytes(), nonce.tobytes(), len(ct)), dtype=np.uint8
+        _keystream(key_np.tobytes(), nonce.tobytes(), len(ct)), dtype=np.uint8
     )
     pt_bytes = (ct ^ stream).astype(np.uint8)
-    return pt_bytes
+    return TensorValue(pt_bytes)
 
 
 @kernel_def("crypto.kem_keygen")
-def _crypto_kem_keygen(pfunc: PFunction) -> Any:
+def _crypto_kem_keygen(pfunc: PFunction) -> tuple[TensorValue, TensorValue]:
     rng = _get_rng()
     sk = rng.integers(0, 256, size=(32,), dtype=np.uint8)
-    pk = np.frombuffer(blake2b(sk.tobytes())[:32], dtype=np.uint8)
-    return (sk, pk)
+    pk_bytes = blake2b(sk.tobytes())[:32]
+    pk = np.frombuffer(pk_bytes, dtype=np.uint8)
+    return (TensorValue(sk), TensorValue(pk))
 
 
 @kernel_def("crypto.kem_derive")
-def _crypto_kem_derive(pfunc: PFunction, sk: Any, peer_pk: Any) -> Any:
-    sk = np.asarray(sk, dtype=np.uint8)
-    peer_pk = np.asarray(peer_pk, dtype=np.uint8)
-    self_pk = np.frombuffer(blake2b(sk.tobytes())[:32], dtype=np.uint8)
-    xored = (self_pk ^ peer_pk).astype(np.uint8)
+def _crypto_kem_derive(
+    pfunc: PFunction, sk: TensorValue, peer_pk: TensorValue
+) -> TensorValue:
+    sk_np = sk.to_numpy().astype(np.uint8, copy=False)
+    peer_pk_np = peer_pk.to_numpy().astype(np.uint8, copy=False)
+
+    self_pk_bytes = blake2b(sk_np.tobytes())[:32]
+    self_pk_arr = np.frombuffer(self_pk_bytes, dtype=np.uint8)
+    xored = (self_pk_arr ^ peer_pk_np).astype(np.uint8)
     secret = np.frombuffer(blake2b(xored.tobytes())[:32], dtype=np.uint8)
-    return secret
+    return TensorValue(secret)
 
 
 @kernel_def("crypto.hkdf")
-def _crypto_hkdf(pfunc: PFunction, secret: Any) -> Any:
-    secret = np.asarray(secret, dtype=np.uint8)
+def _crypto_hkdf(pfunc: PFunction, secret: TensorValue) -> TensorValue:
+    secret_np = secret.to_numpy().astype(np.uint8, copy=False)
     info_str = str(pfunc.attrs.get("info", ""))
     info = info_str.encode("utf-8")
-    out = np.frombuffer(blake2b(secret.tobytes() + info)[:32], dtype=np.uint8)
-    return out
+    out = np.frombuffer(blake2b(secret_np.tobytes() + info)[:32], dtype=np.uint8)
+    return TensorValue(out)