mplang-nightly 0.1.dev164__py3-none-any.whl → 0.1.dev166__py3-none-any.whl

mplang/kernels/phe.py

@@ -14,15 +14,27 @@
 
 """PHE (Partially Homomorphic Encryption) backend implementation using lightPHE."""
 
-from typing import Any
+from __future__ import annotations
+
+import json
+from typing import Any, ClassVar
 
 import numpy as np
 from lightphe import LightPHE
+from lightphe.models.Ciphertext import Ciphertext
 
 from mplang.core.dtype import DType
-from mplang.core.mptype import TensorLike
 from mplang.core.pfunc import PFunction
 from mplang.kernels.base import kernel_def
+from mplang.kernels.value import (
+    TensorValue,
+    Value,
+    ValueDecodeError,
+    ValueProtoBuilder,
+    ValueProtoReader,
+    register_value,
+)
+from mplang.protos.v1alpha1 import value_pb2 as _value_pb2
 
 # This controls the decimal precision used in lightPHE for float operations
 # we force it to 0 to only support integer operations
@@ -30,8 +42,12 @@ from mplang.kernels.base import kernel_def
 PRECISION = 0
 
 
-class PublicKey:
-    """PHE Public Key that implements TensorLike protocol."""
+@register_value
+class PublicKey(Value):
+    """PHE Public Key Value type."""
+
+    KIND: ClassVar[str] = "mplang.phe.PublicKey"
+    WIRE_VERSION: ClassVar[int] = 1
 
     def __init__(
         self,
@@ -62,12 +78,56 @@ class PublicKey:
         """Maximum float value that can be encoded."""
         return float(self.max_value / (2**self.fxp_bits))
 
+    def to_proto(self) -> _value_pb2.ValueProto:
+        """Serialize PublicKey to wire format."""
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_attr("scheme", self.scheme)
+            .set_attr("key_size", self.key_size)
+            .set_attr("max_value", self.max_value)
+            .set_attr("fxp_bits", self.fxp_bits)
+            .set_attr("modulus", str(self.modulus) if self.modulus is not None else "")
+            .set_payload(json.dumps(self.key_data).encode("utf-8"))
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> PublicKey:
+        """Deserialize PublicKey from wire format."""
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported PublicKey version {reader.version}")
+
+        # Read metadata from runtime_attrs
+        scheme = reader.get_attr("scheme")
+        key_size = reader.get_attr("key_size")
+        max_value = reader.get_attr("max_value")
+        fxp_bits = reader.get_attr("fxp_bits")
+        modulus_str = reader.get_attr("modulus")
+        modulus = None if modulus_str == "" else int(modulus_str)
+
+        # JSON deserialize the public key dict
+        key_data = json.loads(reader.payload.decode("utf-8"))
+
+        return cls(
+            key_data=key_data,
+            scheme=scheme,
+            key_size=key_size,
+            max_value=max_value,
+            fxp_bits=fxp_bits,
+            modulus=modulus,
+        )
+
     def __repr__(self) -> str:
         return f"PublicKey(scheme={self.scheme}, key_size={self.key_size}, max_value={self.max_value}, fxp_bits={self.fxp_bits})"
 
 
-class PrivateKey:
-    """PHE Private Key that implements TensorLike protocol."""
+@register_value
+class PrivateKey(Value):
+    """PHE Private Key Value type."""
+
+    KIND: ClassVar[str] = "mplang.phe.PrivateKey"
+    WIRE_VERSION: ClassVar[int] = 1
 
     def __init__(
         self,
@@ -100,12 +160,63 @@ class PrivateKey:
         """Maximum float value that can be encoded."""
         return float(self.max_value / (2**self.fxp_bits))
 
+    def to_proto(self) -> _value_pb2.ValueProto:
+        """Serialize PrivateKey to wire format."""
+        # JSON serialize both key dicts (contain int values)
+        # Store both keys in a single dict to avoid needing length metadata
+        keys_dict = {"sk": self.sk_data, "pk": self.pk_data}
+
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_attr("scheme", self.scheme)
+            .set_attr("key_size", self.key_size)
+            .set_attr("max_value", self.max_value)
+            .set_attr("fxp_bits", self.fxp_bits)
+            .set_attr("modulus", str(self.modulus) if self.modulus is not None else "")
+            .set_payload(json.dumps(keys_dict).encode("utf-8"))
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> PrivateKey:
+        """Deserialize PrivateKey from wire format."""
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported PrivateKey version {reader.version}")
+
+        # Read metadata from runtime_attrs
+        scheme = reader.get_attr("scheme")
+        key_size = reader.get_attr("key_size")
+        max_value = reader.get_attr("max_value")
+        fxp_bits = reader.get_attr("fxp_bits")
+        modulus_str = reader.get_attr("modulus")
+        modulus = None if modulus_str == "" else int(modulus_str)
+
+        # JSON deserialize both key dicts
+        keys_dict = json.loads(reader.payload.decode("utf-8"))
+        sk_data = keys_dict["sk"]
+        pk_data = keys_dict["pk"]
+
+        return cls(
+            sk_data=sk_data,
+            pk_data=pk_data,
+            scheme=scheme,
+            key_size=key_size,
+            max_value=max_value,
+            fxp_bits=fxp_bits,
+            modulus=modulus,
+        )
+
     def __repr__(self) -> str:
         return f"PrivateKey(scheme={self.scheme}, key_size={self.key_size}, max_value={self.max_value}, fxp_bits={self.fxp_bits})"
 
 
-class CipherText:
-    """PHE CipherText that implements TensorLike protocol."""
+@register_value
+class CipherText(Value):
+    """PHE CipherText Value type."""
+
+    KIND: ClassVar[str] = "mplang.phe.CipherText"
+    WIRE_VERSION: ClassVar[int] = 1
 
     def __init__(
         self,
@@ -142,6 +253,106 @@ class CipherText:
         """Maximum float value that can be encoded."""
         return float(self.max_value / (2**self.fxp_bits))
 
+    def to_proto(self) -> _value_pb2.ValueProto:
+        """Serialize CipherText to wire format.
+
+        WARNING: This serialization is tightly coupled to lightphe.Ciphertext
+        internal attributes (value, algorithm_name, keys). Any changes to these
+        attributes in future lightphe versions will break serialization.
+
+        TODO: Check if lightphe provides official serialization methods and
+        migrate to them if available. Consider adding version compatibility checks.
+        """
+        # JSON serialize ciphertext components
+        # ct_data is a list of lightPHE Ciphertext objects
+        # Each Ciphertext has: value, algorithm_name, keys
+        # We need to serialize the list of ciphertexts
+        if not isinstance(self.ct_data, list):
+            raise ValueError(f"ct_data should be a list, got {type(self.ct_data)}")
+
+        ct_list = []
+        for ct in self.ct_data:
+            if not isinstance(ct, Ciphertext):
+                raise TypeError(
+                    f"ct_data must contain lightphe.Ciphertext objects, got {type(ct).__name__}"
+                )
+            ct_list.append({
+                "value": ct.value,
+                "algorithm_name": ct.algorithm_name,
+                "keys": ct.keys,
+            })
+
+        # Combine ct_data and pk_data into single dict
+        payload_dict = {
+            "ct_list": ct_list,
+            "pk": self.pk_data if self.pk_data is not None else None,
+        }
+
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_attr("semantic_dtype", str(self.semantic_dtype))
+            .set_attr("semantic_shape", list(self.semantic_shape))
+            .set_attr("scheme", self.scheme)
+            .set_attr("key_size", self.key_size)
+            .set_attr("max_value", self.max_value)
+            .set_attr("fxp_bits", self.fxp_bits)
+            .set_attr("modulus", str(self.modulus) if self.modulus is not None else "")
+            .set_payload(json.dumps(payload_dict).encode("utf-8"))
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> CipherText:
+        """Deserialize CipherText from wire format."""
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported CipherText version {reader.version}")
+
+        # Read metadata from runtime_attrs
+        semantic_dtype_str = reader.get_attr("semantic_dtype")
+        semantic_shape = reader.get_attr("semantic_shape")
+        scheme = reader.get_attr("scheme")
+        key_size = reader.get_attr("key_size")
+        max_value = reader.get_attr("max_value")
+        fxp_bits = reader.get_attr("fxp_bits")
+        modulus_str = reader.get_attr("modulus")
+        modulus = None if modulus_str == "" else int(modulus_str)
+
+        # JSON deserialize ciphertext and public key
+        payload_dict = json.loads(reader.payload.decode("utf-8"))
+        ct_list = payload_dict["ct_list"]
+        pk_data = payload_dict["pk"]
+
+        # Reconstruct ct_data: list of Ciphertext objects
+        ct_data = []
+        for ct_dict in ct_list:
+            if ct_dict["keys"] is None or ct_dict["algorithm_name"] is None:
+                raise ValueDecodeError(
+                    "Invalid CipherText: missing keys or algorithm_name in serialized data"
+                )
+            ct_data.append(
+                Ciphertext(
+                    algorithm_name=ct_dict["algorithm_name"],
+                    keys=ct_dict["keys"],
+                    value=ct_dict["value"],
+                )
+            )
+
+        # Parse dtype string back to DType
+        dtype = DType.from_any(semantic_dtype_str)
+
+        return cls(
+            ct_data=ct_data,
+            semantic_dtype=dtype,
+            semantic_shape=tuple(semantic_shape),
+            scheme=scheme,
+            key_size=key_size,
+            pk_data=pk_data,
+            max_value=max_value,
+            fxp_bits=fxp_bits,
+            modulus=modulus,
+        )
+
     def __repr__(self) -> str:
         return f"CipherText(dtype={self.semantic_dtype}, shape={self.semantic_shape}, scheme={self.scheme})"
 
@@ -257,23 +468,6 @@ def _range_decode_mixed(
         return _range_decode_integer(encoded_value, max_value, modulus)
 
 
-def _convert_to_numpy(obj: TensorLike) -> np.ndarray:
-    """Convert a TensorLike object to numpy array."""
-    if isinstance(obj, np.ndarray):
-        return obj
-
-    # Try to use .numpy() method if available
-    if hasattr(obj, "numpy"):
-        numpy_method = getattr(obj, "numpy", None)
-        if callable(numpy_method):
-            try:
-                return np.asarray(numpy_method())
-            except Exception:
-                pass
-
-    return np.asarray(obj)
-
-
 @kernel_def("phe.keygen")
 def _phe_keygen(pfunc: PFunction) -> Any:
     scheme = pfunc.attrs.get("scheme", "paillier")
@@ -334,14 +528,16 @@ def _phe_keygen(pfunc: PFunction) -> Any:
 
 
 @kernel_def("phe.encrypt")
-def _phe_encrypt(pfunc: PFunction, plaintext: Any, public_key: PublicKey) -> Any:
+def _phe_encrypt(
+    pfunc: PFunction, plaintext: TensorValue, public_key: PublicKey
+) -> Any:
     # Validate public_key type
     if not isinstance(public_key, PublicKey):
         raise ValueError("Second argument must be a PublicKey instance")
 
     try:
         # Convert plaintext to numpy to get semantic type info
-        plaintext_np = _convert_to_numpy(plaintext)
+        plaintext_np = plaintext.to_numpy()
         semantic_dtype = DType.from_numpy(plaintext_np.dtype)
         semantic_shape = plaintext_np.shape
 
@@ -403,14 +599,14 @@ def _phe_encrypt(pfunc: PFunction, plaintext: Any, public_key: PublicKey) -> Any
 
 
 @kernel_def("phe.mul")
-def _phe_mul(pfunc: PFunction, ciphertext: CipherText, plaintext: Any) -> Any:
+def _phe_mul(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorValue) -> Any:
     # Validate that first argument is a CipherText
     if not isinstance(ciphertext, CipherText):
         raise ValueError("First argument must be a CipherText instance")
 
     try:
         # Convert plaintext to numpy
-        plaintext_np = _convert_to_numpy(plaintext)
+        plaintext_np = plaintext.to_numpy()
 
         # Check if plaintext is floating point type - multiplication not supported
         if np.issubdtype(plaintext_np.dtype, np.floating):
@@ -511,7 +707,7 @@ def _phe_add(pfunc: PFunction, lhs: Any, rhs: Any) -> Any:
         elif isinstance(rhs, CipherText):
             return _phe_add_ct2pt(rhs, lhs)
         else:
-            return _convert_to_numpy(lhs) + _convert_to_numpy(rhs)
+            return TensorValue(lhs.to_numpy() + rhs.to_numpy())
     except ValueError:
         raise
     except Exception as e:  # pragma: no cover
@@ -593,9 +789,9 @@ def _phe_add_ct2ct(ct1: CipherText, ct2: CipherText) -> CipherText:
     )
 
 
-def _phe_add_ct2pt(ciphertext: CipherText, plaintext: TensorLike) -> CipherText:
+def _phe_add_ct2pt(ciphertext: CipherText, plaintext: TensorValue) -> CipherText:
     # Convert plaintext to numpy
-    plaintext_np = _convert_to_numpy(plaintext)
+    plaintext_np = plaintext.to_numpy()
     plaintext_dtype = DType.from_numpy(plaintext_np.dtype)
 
     # Check for mixed precision issue: floating point ciphertext + integer plaintext
@@ -816,14 +1012,14 @@ def _phe_decrypt(
             ciphertext.semantic_shape
         )
 
-        return [plaintext_np]
+        return [TensorValue(plaintext_np)]
 
     except Exception as e:
         raise RuntimeError(f"Failed to decrypt data: {e}") from e
 
 
 @kernel_def("phe.dot")
-def _phe_dot(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorLike) -> Any:
+def _phe_dot(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorValue) -> Any:
     """Execute homomorphic dot product with zero-value optimization.
 
     Supports various dot product operations:
@@ -844,7 +1040,7 @@ def _phe_dot(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorLike) ->
 
     try:
         # Convert plaintext to numpy
-        plaintext_np = _convert_to_numpy(plaintext)
+        plaintext_np = plaintext.to_numpy()
 
         # Check if plaintext is floating point type - dot product not supported
         if np.issubdtype(plaintext_np.dtype, np.floating):
@@ -1109,7 +1305,7 @@ def _phe_dot(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorLike) ->
 
 
 @kernel_def("phe.gather")
-def _phe_gather(pfunc: PFunction, ciphertext: CipherText, indices: Any) -> Any:
+def _phe_gather(pfunc: PFunction, ciphertext: CipherText, indices: TensorValue) -> Any:
     """Execute gather operation on CipherText.
 
     Supports gathering from multidimensional CipherText using multidimensional indices.
@@ -1126,7 +1322,7 @@ def _phe_gather(pfunc: PFunction, ciphertext: CipherText, indices: Any) -> Any:
 
     try:
         # Convert indices to numpy
-        indices_np = _convert_to_numpy(indices)
+        indices_np = indices.to_numpy()
 
         if not np.issubdtype(indices_np.dtype, np.integer):
             raise ValueError("Indices must be of integer type")
@@ -1224,7 +1420,10 @@ def _phe_gather(pfunc: PFunction, ciphertext: CipherText, indices: Any) -> Any:
 
 @kernel_def("phe.scatter")
 def _phe_scatter(
-    pfunc: PFunction, ciphertext: CipherText, indices: TensorLike, updated: CipherText
+    pfunc: PFunction,
+    ciphertext: CipherText,
+    indices: TensorValue,
+    updated: CipherText,
 ) -> Any:
     """Execute scatter operation on CipherText.
 
@@ -1252,7 +1451,7 @@ def _phe_scatter(
 
     try:
         # Convert indices to numpy
-        indices_np = _convert_to_numpy(indices)
+        indices_np = indices.to_numpy()
 
         if not np.issubdtype(indices_np.dtype, np.integer):
             raise ValueError("Indices must be of integer type")

mplang/kernels/spu.py

@@ -15,15 +15,37 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import spu.api as spu_api
 import spu.libspu as libspu
 
-from mplang.core.mptype import TensorLike
+from mplang.core.dtype import (
+    BOOL,
+    FLOAT32,
+    FLOAT64,
+    INT8,
+    INT16,
+    INT32,
+    INT64,
+    UINT8,
+    UINT16,
+    UINT32,
+    UINT64,
+    DType,
+)
 from mplang.core.pfunc import PFunction
 from mplang.kernels.base import cur_kctx, kernel_def
+from mplang.kernels.value import (
+    TensorValue,
+    Value,
+    ValueDecodeError,
+    ValueProtoBuilder,
+    ValueProtoReader,
+    register_value,
+)
+from mplang.protos.v1alpha1 import value_pb2 as _value_pb2
 from mplang.runtime.link_comm import LinkCommunicator
 
 
@@ -32,36 +54,106 @@ def shape_spu_to_np(spu_shape: Any) -> tuple[int, ...]:
     return tuple(spu_shape.dims)
 
 
-def dtype_spu_to_np(spu_dtype: Any) -> np.dtype:
-    """Convert SPU dtype to numpy dtype."""
+def dtype_spu_to_mpl(spu_dtype: libspu.DataType) -> DType:
+    """Convert libspu.DataType to MPLang DType."""
     MAP = {
-        libspu.DataType.DT_F32: np.float32,
-        libspu.DataType.DT_F64: np.float64,
-        libspu.DataType.DT_I1: np.bool_,
-        libspu.DataType.DT_I8: np.int8,
-        libspu.DataType.DT_U8: np.uint8,
-        libspu.DataType.DT_I16: np.int16,
-        libspu.DataType.DT_U16: np.uint16,
-        libspu.DataType.DT_I32: np.int32,
-        libspu.DataType.DT_U32: np.uint32,
-        libspu.DataType.DT_I64: np.int64,
-        libspu.DataType.DT_U64: np.uint64,
+        libspu.DataType.DT_F32: FLOAT32,
+        libspu.DataType.DT_F64: FLOAT64,
+        libspu.DataType.DT_I1: BOOL,
+        libspu.DataType.DT_I8: INT8,
+        libspu.DataType.DT_U8: UINT8,
+        libspu.DataType.DT_I16: INT16,
+        libspu.DataType.DT_U16: UINT16,
+        libspu.DataType.DT_I32: INT32,
+        libspu.DataType.DT_U32: UINT32,
+        libspu.DataType.DT_I64: INT64,
+        libspu.DataType.DT_U64: UINT64,
     }
-    return MAP[spu_dtype]  # type: ignore[return-value]
+    return MAP[spu_dtype]
 
 
+@register_value
 @dataclass
-class SpuValue:
-    """SPU value container for secure computation."""
+class SpuValue(Value):
+    """SPU value container for secure computation (Value type)."""
+
+    KIND: ClassVar[str] = "mplang.spu.SpuValue"
+    WIRE_VERSION: ClassVar[int] = 1
 
     shape: tuple[int, ...]
-    dtype: Any
+    dtype: DType  # Now uses MPLang's unified DType
     vtype: libspu.Visibility
     share: libspu.Share
 
     def __repr__(self) -> str:
         return f"SpuValue({self.shape},{self.dtype},{self.vtype})"
 
+    def to_proto(self) -> _value_pb2.ValueProto:
+        """Serialize SpuValue to wire format.
+
+        libspu.Share has two attributes:
+        - meta: bytes (protobuf serialized metadata)
+        - share_chunks: list[bytes] (the actual secret share data)
+
+        Strategy: Store shape/dtype/vtype in runtime_attrs, concatenate share.meta + all chunks in payload.
+        """
+        # Store metadata in runtime_attrs; keep chunk lengths for payload splitting
+        chunk_lengths = [len(chunk) for chunk in self.share.share_chunks]
+
+        # Payload contains only share chunks (meta stored in attrs)
+        payload = b""
+        for chunk in self.share.share_chunks:
+            payload += chunk
+
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_attr("shape", list(self.shape))
+            .set_attr("dtype", self.dtype.name)  # Serialize DType name
+            .set_attr("vtype", int(self.vtype))
+            .set_attr("share_meta", self.share.meta)
+            .set_attr("chunk_lengths", chunk_lengths)
+            .set_payload(payload)
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> SpuValue:
+        """Deserialize SpuValue from wire format."""
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported SpuValue version {reader.version}")
+
+        # Read metadata from runtime_attrs
+        shape = tuple(reader.get_attr("shape"))
+        dtype_name = reader.get_attr("dtype")
+        # Reconstruct DType from serialized name (numpy dtype string)
+        dtype = DType.from_numpy(dtype_name)
+        vtype = libspu.Visibility(reader.get_attr("vtype"))
+        share_meta = reader.get_attr("share_meta")
+        chunk_lengths = reader.get_attr("chunk_lengths")
+
+        # Parse payload: [chunk_0][chunk_1]...
+        payload = reader.payload
+        offset = 0
+
+        share_chunks: list[bytes] = []
+        for chunk_len in chunk_lengths:
+            chunk = payload[offset : offset + chunk_len]
+            offset += chunk_len
+            share_chunks.append(chunk)
+
+        # Reconstruct libspu.Share
+        share = libspu.Share()
+        share.meta = share_meta
+        share.share_chunks = share_chunks
+
+        return cls(
+            shape=shape,
+            dtype=dtype,
+            vtype=vtype,
+            share=share,
+        )
+
 
 def _get_spu_config_and_world() -> tuple[libspu.RuntimeConfig, int]:
     kctx = cur_kctx()
@@ -128,33 +220,25 @@ def _spu_seed_env(pfunc: PFunction, *args: Any) -> Any:
 
 
 @kernel_def("spu.makeshares")
-def _spu_makeshares(pfunc: PFunction, *args: Any) -> Any:
-    """Create SPU shares from input data.
-
-    Args:
-        pfunc: PFunction containing makeshares metadata
-        args: Input data to be shared (single tensor)
-
-    Returns:
-        Tuple of SPU shares (SpuValue), one for each party.
-    """
-    assert len(args) == 1
-
+def _spu_makeshares(pfunc: PFunction, tensor: TensorValue) -> tuple[SpuValue, ...]:
+    """Create SPU shares from input TensorValue data."""
     visibility_value = pfunc.attrs.get("visibility", libspu.Visibility.VIS_SECRET.value)
     if isinstance(visibility_value, int):
         visibility = libspu.Visibility(visibility_value)
     else:
         visibility = visibility_value
 
-    arg = np.array(args[0], copy=False)
+    arg = tensor.to_numpy()
     cfg, world = _get_spu_config_and_world()
     spu_io = spu_api.Io(world, cfg)
     shares = spu_io.make_shares(arg, visibility)
     assert len(shares) == world, f"Expected {world} shares, got {len(shares)}"
+    # Store MPLang DType instead of libspu.DataType
+    dtype = DType.from_numpy(arg.dtype)
     return tuple(
         SpuValue(
             shape=arg.shape,
-            dtype=arg.dtype,
+            dtype=dtype,
             vtype=visibility,
             share=share,
         )
@@ -163,24 +247,29 @@ def _spu_makeshares(pfunc: PFunction, *args: Any) -> Any:
 
 
 @kernel_def("spu.reconstruct")
-def _spu_reconstruct(pfunc: PFunction, *args: Any) -> Any:
+def _spu_reconstruct(pfunc: PFunction, *shares: SpuValue) -> TensorValue:
     """Reconstruct plaintext data from SPU shares."""
     cfg, world = _get_spu_config_and_world()
-    assert len(args) == world, f"Expected {world} shares, got {len(args)}"
-    for i, arg in enumerate(args):
-        if not isinstance(arg, SpuValue):
+    assert len(shares) == world, f"Expected {world} shares, got {len(shares)}"
+    for i, share in enumerate(shares):
+        if not isinstance(share, SpuValue):
             raise ValueError(
-                f"Input {i} must be SpuValue, got {type(arg)}. Reconstruction requires SPU shares as input."
+                f"Input {i} must be SpuValue, got {type(share)}. Reconstruction requires SPU shares as input."
             )
-    spu_args: list[SpuValue] = list(args)  # type: ignore
-    shares = [spu_arg.share for spu_arg in spu_args]
+    spu_args: list[SpuValue] = list(shares)  # type: ignore
+    share_payloads = [spu_arg.share for spu_arg in spu_args]
     spu_io = spu_api.Io(world, cfg)
-    reconstructed = spu_io.reconstruct(shares)
-    return reconstructed
+    reconstructed = spu_io.reconstruct(share_payloads)
+    base = np.array(reconstructed, copy=False)
+    # Respect semantic dtype/shape recorded on shares (all shares share same meta).
+    semantic_dtype = shares[0].dtype.to_numpy()  # DType now has to_numpy() method
+    semantic_shape = shares[0].shape
+    restored = np.asarray(base, dtype=semantic_dtype).reshape(semantic_shape)
+    return TensorValue(np.array(restored, copy=False))
 
 
 @kernel_def("spu.run_pphlo")
-def _spu_run_mlir(pfunc: PFunction, *args: Any) -> Any:
+def _spu_run_mlir(pfunc: PFunction, *args: SpuValue) -> tuple[SpuValue, ...]:
     """Execute compiled SPU function (spu.run_pphlo) and return SpuValue outputs.
 
     Participation rule: a rank participates iff its entry in the stored
@@ -240,10 +329,10 @@ def _spu_run_mlir(pfunc: PFunction, *args: Any) -> Any:
     spu_rt.run(executable)
     shares = [spu_rt.get_var(out_name) for out_name in output_names]
     metas = [spu_rt.get_var_meta(out_name) for out_name in output_names]
-    results: list[TensorLike] = [
+    results: list[SpuValue] = [
         SpuValue(
             shape=shape_spu_to_np(meta.shape),
-            dtype=dtype_spu_to_np(meta.data_type),
+            dtype=dtype_spu_to_mpl(meta.data_type),
             vtype=meta.visibility,
             share=shares[idx],
         )