mplang-nightly 0.1.dev192__py3-none-any.whl → 0.1.dev268__py3-none-any.whl

mplang/{kernels → v1/kernels}/basic.py

@@ -16,11 +16,11 @@ from __future__ import annotations
 
 import numpy as np
 
-from mplang.core import PFunction, TableType, 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
+from mplang.v1.core import PFunction, TableType, TensorType
+from mplang.v1.kernels.base import cur_kctx, kernel_def
+from mplang.v1.kernels.value import TableValue, TensorValue, Value
+from mplang.v1.runtime.data_providers import get_provider, resolve_uri
+from mplang.v1.utils import table_utils
 
 
 @kernel_def("basic.identity")
@@ -45,17 +45,17 @@ def _read(pfunc: PFunction) -> Value:
     except Exception as e:  # pragma: no cover - provider errors
         raise RuntimeError(f"basic.read failed: {e}") from e
 
+    if isinstance(data, Value):
+        return data
+
     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
+    elif isinstance(out_t, TensorType):
         return TensorValue(np.asarray(data))
-    raise TypeError(
-        f"basic.read only supports TableType/TensorType outputs, got {type(out_t).__name__}"
-    )
+    else:
+        raise TypeError(
+            f"basic.read only supports TableType/TensorType outputs, got {type(out_t).__name__}"
+        )
 
 
 @kernel_def("basic.write")
@@ -85,9 +85,9 @@ def _constant(pfunc: PFunction) -> Value:
     out_t = pfunc.outs_info[0]
     fmt = pfunc.attrs.get("data_format")
     if isinstance(out_t, TableType):
-        if fmt != "bytes[csv]":
+        if fmt != "bytes[parquet]":
             raise ValueError(f"unsupported table constant format {fmt}")
-        df = table_utils.csv_to_dataframe(data_bytes)
+        df = table_utils.decode_table(data_bytes, format="parquet")
         return TableValue(df)
     # tensor path
     shape = out_t.shape  # type: ignore[attr-defined,union-attr]

mplang/{kernels → v1/kernels}/context.py

@@ -17,12 +17,12 @@ from __future__ import annotations
 from collections.abc import Mapping
 from typing import Any
 
-from mplang.core.dtypes import UINT8, DType
-from mplang.core.pfunc import PFunction
-from mplang.core.table import TableLike, TableType
-from mplang.core.tensor import TensorLike, TensorType
-from mplang.kernels import base
-from mplang.kernels.base import KernelContext, get_kernel_spec, kernel_exists
+from mplang.v1.core.dtypes import UINT8, DType
+from mplang.v1.core.pfunc import PFunction
+from mplang.v1.core.table import PandasTableLike, TableLike, TableType
+from mplang.v1.core.tensor import TensorLike, TensorType
+from mplang.v1.kernels import base
+from mplang.v1.kernels.base import KernelContext, get_kernel_spec, kernel_exists
 
 # Default bindings
 # Import kernel implementation modules explicitly so their @kernel_def entries
@@ -35,14 +35,14 @@ def _ensure_impl_imported() -> None:
     global _IMPL_IMPORTED
     if _IMPL_IMPORTED:
         return
-    from mplang.kernels import basic as _impl_basic  # noqa: F401
-    from mplang.kernels import crypto as _impl_crypto  # noqa: F401
-    from mplang.kernels import fhe as _impl_fhe  # noqa: F401
-    from mplang.kernels import mock_tee as _impl_tee  # noqa: F401
-    from mplang.kernels import phe as _impl_phe  # noqa: F401
-    from mplang.kernels import spu as _impl_spu  # noqa: F401
-    from mplang.kernels import sql_duckdb as _impl_sql_duckdb  # noqa: F401
-    from mplang.kernels import stablehlo as _impl_stablehlo  # noqa: F401
+    from mplang.v1.kernels import basic as _impl_basic  # noqa: F401
+    from mplang.v1.kernels import crypto as _impl_crypto  # noqa: F401
+    from mplang.v1.kernels import fhe as _impl_fhe  # noqa: F401
+    from mplang.v1.kernels import mock_tee as _impl_tee  # noqa: F401
+    from mplang.v1.kernels import phe as _impl_phe  # noqa: F401
+    from mplang.v1.kernels import spu as _impl_spu  # noqa: F401
+    from mplang.v1.kernels import sql_duckdb as _impl_sql_duckdb  # noqa: F401
+    from mplang.v1.kernels import stablehlo as _impl_stablehlo  # noqa: F401
 
     _IMPL_IMPORTED = True
 
@@ -317,9 +317,12 @@ def _validate_table_arg(
         raise TypeError(
             f"kernel {fn_type} input[{arg_index}] expects TableLike, got {type(value).__name__}"
         )
-    if len(value.columns) != len(spec.columns):
+    columns = (
+        value.columns if isinstance(value, PandasTableLike) else value.column_names
+    )
+    if len(columns) != len(spec.columns):
         raise ValueError(
-            f"kernel {fn_type} input[{arg_index}] column count mismatch: got {len(value.columns)}, expected {len(spec.columns)}"
+            f"kernel {fn_type} input[{arg_index}] column count mismatch: got {len(columns)}, expected {len(spec.columns)}"
         )
 
 

mplang/{kernels → v1/kernels}/crypto.py

@@ -18,10 +18,10 @@ import os
 
 import numpy as np
 
-from mplang.core import PFunction
-from mplang.kernels.base import cur_kctx, kernel_def
-from mplang.kernels.value import TensorValue
-from mplang.utils.crypto import blake2b
+from mplang.v1.core import PFunction
+from mplang.v1.kernels.base import cur_kctx, kernel_def
+from mplang.v1.kernels.value import TensorValue
+from mplang.v1.utils.crypto import blake2b
 
 __all__: list[str] = []  # No public exports currently
 
@@ -45,11 +45,9 @@ def _get_rng() -> np.random.Generator:
 def _keystream(key: bytes, nonce: bytes, length: int) -> bytes:
     # WARNING (INSECURE): hash-based keystream (key||nonce||counter)
     out = bytearray()
-    counter = 0
     while len(out) < length:
-        chunk = blake2b(key + nonce + counter.to_bytes(4, "little"))
+        chunk = blake2b(key + nonce)
         out.extend(chunk)
-        counter += 1
     return bytes(out[:length])
 
 
@@ -68,7 +66,7 @@ def _crypto_encrypt(
     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)
+    nonce = rng.integers(0, 256, size=(16,), dtype=np.uint8)
     stream = np.frombuffer(
         _keystream(key_np.tobytes(), nonce.tobytes(), pt_bytes_np.size), dtype=np.uint8
     )
@@ -83,8 +81,8 @@ def _crypto_decrypt(
 ) -> 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:]
+    nonce = ct_np[:16]
+    ct = ct_np[16:]
     stream = np.frombuffer(
         _keystream(key_np.tobytes(), nonce.tobytes(), len(ct)), dtype=np.uint8
     )

mplang/{kernels → v1/kernels}/fhe.py

@@ -23,8 +23,9 @@ from typing import Any
 import numpy as np
 import tenseal as ts
 
-from mplang.core import DType, PFunction, TensorLike
-from mplang.kernels.base import kernel_def
+from mplang.v1.core import DType, PFunction, TensorLike
+from mplang.v1.kernels.base import kernel_def
+from mplang.v1.kernels.value import TensorValue
 
 
 class FHEContext:
@@ -337,13 +338,14 @@ def _fhe_decrypt(pfunc: PFunction, ciphertext: CipherText, context: FHEContext)
 
         # Restore original shape
         if ciphertext.semantic_shape == ():
-            # Was a scalar, extract single value
-            result = decrypted_np[0:1].reshape(())
+            # Scalar: shape ()
+            result_np = decrypted_np[0:1].reshape(())
         else:
-            # Keep as vector
-            result = decrypted_np
+            # Vector: keep 1D array
+            result_np = decrypted_np
 
-        return (result,)
+        # Return TensorValue to adhere to kernel Value I/O convention
+        return (TensorValue(np.asarray(result_np)),)
 
     except Exception as e:
         raise RuntimeError(f"FHE vector decryption failed: {e}") from e

mplang/{kernels → v1/kernels}/mock_tee.py

@@ -20,9 +20,9 @@ import warnings
 import numpy as np
 from numpy.typing import NDArray
 
-from mplang.core import PFunction
-from mplang.kernels.base import cur_kctx, kernel_def
-from mplang.kernels.value import TensorValue
+from mplang.v1.core import PFunction
+from mplang.v1.kernels.base import cur_kctx, kernel_def
+from mplang.v1.kernels.value import TensorValue
 
 __all__: list[str] = []
 

mplang/{kernels → v1/kernels}/phe.py

@@ -23,9 +23,9 @@ import numpy as np
 from lightphe import LightPHE
 from lightphe.models.Ciphertext import Ciphertext
 
-from mplang.core import DType, PFunction
-from mplang.kernels.base import kernel_def
-from mplang.kernels.value import (
+from mplang.v1.core import DType, PFunction
+from mplang.v1.kernels.base import kernel_def
+from mplang.v1.kernels.value import (
     TensorValue,
     Value,
     ValueDecodeError,
@@ -33,7 +33,7 @@ from mplang.kernels.value import (
     ValueProtoReader,
     register_value,
 )
-from mplang.protos.v1alpha1 import value_pb2 as _value_pb2
+from mplang.v1.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
@@ -473,10 +473,9 @@ def _phe_keygen(pfunc: PFunction) -> Any:
     # use small key_size to speed up tests
     # in production use at least 2048 bits or 3072 bits for better security
     key_size = pfunc.attrs.get("key_size", 2048)
-    max_value = pfunc.attrs.get(
-        "max_value", 2**32
-    )  # Use larger range to avoid overflow
-    fxp_bits = pfunc.attrs.get("fxp_bits", 12)
+    # Accept very large max_value; allow decimal string input, kept simple like other attrs
+    max_value = int(pfunc.attrs.get("max_value", 2**32))
+    fxp_bits = int(pfunc.attrs.get("fxp_bits", 12))
 
     # Validate scheme
     if scheme.lower() not in ["paillier"]:
@@ -638,7 +637,8 @@ def _phe_mul(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorValue) -
             # Use numpy to create a properly broadcasted index mapping
             # Create a dummy array with same shape as ciphertext, fill with indices
             dummy_ct = (
-                np.arange(np.prod(ciphertext.semantic_shape))
+                np
+                .arange(np.prod(ciphertext.semantic_shape))
                 .reshape(ciphertext.semantic_shape)
                 .astype(np.int64)
             )
@@ -745,7 +745,8 @@ def _phe_add_ct2ct(ct1: CipherText, ct2: CipherText) -> CipherText:
     # Broadcast ct1 if needed
     if ct1.semantic_shape != result_shape:
         dummy_ct1 = (
-            np.arange(np.prod(ct1.semantic_shape))
+            np
+            .arange(np.prod(ct1.semantic_shape))
             .reshape(ct1.semantic_shape)
             .astype(np.int64)
         )
@@ -758,7 +759,8 @@ def _phe_add_ct2ct(ct1: CipherText, ct2: CipherText) -> CipherText:
     # Broadcast ct2 if needed
     if ct2.semantic_shape != result_shape:
         dummy_ct2 = (
-            np.arange(np.prod(ct2.semantic_shape))
+            np
+            .arange(np.prod(ct2.semantic_shape))
             .reshape(ct2.semantic_shape)
             .astype(np.int64)
         )
@@ -831,7 +833,8 @@ def _phe_add_ct2pt(ciphertext: CipherText, plaintext: TensorValue) -> CipherText
     # Broadcast ciphertext if needed
     if ciphertext.semantic_shape != result_shape:
         dummy_ct = (
-            np.arange(np.prod(ciphertext.semantic_shape))
+            np
+            .arange(np.prod(ciphertext.semantic_shape))
             .reshape(ciphertext.semantic_shape)
             .astype(np.int64)
         )
@@ -997,12 +1000,17 @@ def _phe_decrypt(
         # Convert to target dtype
         if target_dtype.kind in "iu":  # integer types
             # Convert floats back to integers for integer semantic types
-            processed_data = [round(val) for val in decoded_data]
-            # Handle overflow for smaller integer types
-            info = np.iinfo(target_dtype)
-            processed_data = [
-                max(info.min, min(info.max, val)) for val in processed_data
-            ]
+            # decoded_data are numeric (ints or floats); normalize to Python int
+            ints = [round(v) if isinstance(v, float) else v for v in decoded_data]
+            if np.issubdtype(target_dtype, np.unsignedinteger):
+                # Reduce modulo 2^k for unsigned to preserve ring semantics
+                width = np.iinfo(target_dtype).bits
+                mod = 1 << width
+                processed_data = [v % mod for v in ints]
+            else:
+                # Signed integers: clamp to dtype range
+                info = np.iinfo(target_dtype)
+                processed_data = [max(info.min, min(info.max, v)) for v in ints]
         else:  # float types
             processed_data = decoded_data
 

mplang/{kernels → v1/kernels}/spu.py

@@ -21,7 +21,7 @@ import numpy as np
 import spu.api as spu_api
 import spu.libspu as libspu
 
-from mplang.core import (
+from mplang.v1.core import (
     BOOL,
     FLOAT32,
     FLOAT64,
@@ -36,8 +36,8 @@ from mplang.core import (
     DType,
     PFunction,
 )
-from mplang.kernels.base import cur_kctx, kernel_def
-from mplang.kernels.value import (
+from mplang.v1.kernels.base import cur_kctx, kernel_def
+from mplang.v1.kernels.value import (
     TensorValue,
     Value,
     ValueDecodeError,
@@ -45,8 +45,8 @@ from mplang.kernels.value import (
     ValueProtoReader,
     register_value,
 )
-from mplang.protos.v1alpha1 import value_pb2 as _value_pb2
-from mplang.runtime.link_comm import LinkCommunicator
+from mplang.v1.protos.v1alpha1 import value_pb2 as _value_pb2
+from mplang.v1.runtime.link_comm import LinkCommunicator
 
 
 def shape_spu_to_np(spu_shape: Any) -> tuple[int, ...]:

mplang/{kernels → v1/kernels}/sql_duckdb.py

@@ -14,9 +14,9 @@
 
 from __future__ import annotations
 
-from mplang.core import PFunction
-from mplang.kernels.base import kernel_def
-from mplang.kernels.value import TableValue
+from mplang.v1.core import PFunction
+from mplang.v1.kernels.base import kernel_def
+from mplang.v1.kernels.value import TableValue
 
 
 @kernel_def("duckdb.run_sql")
@@ -38,5 +38,7 @@ def _duckdb_sql(pfunc: PFunction, *args: TableValue) -> TableValue:
             arrow_table = arg.to_arrow()
             conn.register(name, arrow_table)
     # Fetch result as Arrow table for consistency
+    if pfunc.fn_text is None:
+        raise ValueError("SQL function text is None")
     res_arrow = conn.execute(pfunc.fn_text).fetch_arrow_table()
     return TableValue(res_arrow)

mplang/{kernels → v1/kernels}/stablehlo.py

@@ -17,14 +17,14 @@ from __future__ import annotations
 from typing import Any
 
 import jax
+import jax.extend as jxt
 import jax.numpy as jnp
 import numpy as np
-from jax._src import xla_bridge
-from jax.lib import xla_client as xc
+from jax._src import compiler
 
-from mplang.core import PFunction
-from mplang.kernels.base import cur_kctx, kernel_def
-from mplang.kernels.value import TensorValue
+from mplang.v1.core import PFunction
+from mplang.v1.kernels.base import cur_kctx, kernel_def
+from mplang.v1.kernels.value import TensorValue
 
 
 @kernel_def("mlir.stablehlo")
@@ -47,11 +47,13 @@ def _stablehlo_exec(pfunc: PFunction, *args: Any) -> Any:
     key = f"stablehlo.compile_cache.{h}"
     compiled = rt.get_state(key)
     if compiled is None:
-        backend = jax.default_backend()
-        client = xla_bridge.get_backend(backend)
-        compile_options = xc.CompileOptions()
+        client = jxt.backend.get_backend()
+        compile_options = compiler.get_compile_options(num_replicas=1, num_partitions=1)
+
         try:
-            compiled = client.compile(mlir_text, compile_options)
+            compiled = client.compile_and_load(
+                mlir_text, client.devices(), compile_options
+            )
         except Exception as e:  # pragma: no cover
             raise RuntimeError(f"StableHLO compile failed: {e}") from e
         rt.set_state(key, compiled)
@@ -76,14 +78,13 @@ def _stablehlo_exec(pfunc: PFunction, *args: Any) -> Any:
     ]
 
     try:
-        result = compiled.execute_sharded(jax_args)
-        arrays = result.disassemble_into_single_device_arrays()
-        flat: list[Any] = []
-        for lst in arrays:
-            if isinstance(lst, list) and len(lst) == 1:
-                flat.append(TensorValue(np.asarray(lst[0])))
-            else:
-                flat.extend(TensorValue(np.asarray(a)) for a in lst)
+        # Execute with the new LoadedExecutable interface
+        result = compiled.execute(jax_args)
+
+        # Use jax.tree_util.tree_flatten to robustly handle any PyTree structure
+        flat_results, _ = jax.tree_util.tree_flatten(result)
+        flat = [TensorValue(np.asarray(item)) for item in flat_results]
+
         return tuple(flat)
     except Exception as e:  # pragma: no cover
         raise RuntimeError(f"StableHLO execute failed: {e}") from e

mplang/{kernels → v1/kernels}/value.py

@@ -17,7 +17,7 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, ClassVar, TypeVar
 
-from mplang.protos.v1alpha1 import value_pb2 as _value_pb2
+from mplang.v1.protos.v1alpha1 import value_pb2 as _value_pb2
 
 if TYPE_CHECKING:
     import numpy as np
@@ -591,7 +591,7 @@ class TableValue(Value):  # well-known table (Arrow IPC) Value
 
         Note: This creates a copy and converts from Arrow to pandas format.
         For better performance, consider using to_arrow() and working with
-        Arrow-native APIs (DuckDB, Ibis, etc.) directly.
+        Arrow-native APIs (DuckDB, etc.) directly.
 
         Returns:
             pandas.DataFrame: Converted dataframe

mplang/{ops → v1/ops}/__init__.py

@@ -19,15 +19,15 @@ This module contains compilers that transform high-level functions into
 portable, serializable intermediate representations.
 """
 
-from mplang.ops import basic, crypto, ibis_cc, jax_cc, phe, spu, sql_cc, tee
-from mplang.ops.base import FeOperation as FeOperation
+from mplang.v1.ops import basic, crypto, jax_cc, nnx_cc, phe, spu, sql_cc, tee
+from mplang.v1.ops.base import FeOperation as FeOperation
 
 __all__ = [
     "FeOperation",
     "basic",
     "crypto",
-    "ibis_cc",
     "jax_cc",
+    "nnx_cc",
     "phe",
     "spu",
     "sql_cc",

mplang/{ops → v1/ops}/base.py

@@ -20,7 +20,7 @@ from typing import Any
 
 from jax.tree_util import PyTreeDef, tree_flatten
 
-from mplang.core import MPContext, MPObject, PFunction, TableType, TensorType
+from mplang.v1.core import MPContext, MPObject, PFunction, TableType, TensorType
 
 # -----------------------------------------------------------------------------
 # Triad ABI

mplang/{ops → v1/ops}/basic.py

@@ -15,7 +15,7 @@
 
 from jax.tree_util import PyTreeDef, tree_flatten
 
-from mplang.core import (
+from mplang.v1.core import (
     UINT8,
     UINT64,
     MPObject,
@@ -27,8 +27,8 @@ from mplang.core import (
     TensorLike,
     TensorType,
 )
-from mplang.ops.base import stateless_mod
-from mplang.utils import table_utils
+from mplang.v1.ops.base import stateless_mod
+from mplang.v1.utils import table_utils
 
 _BASIC_MOD = stateless_mod("basic")
 
@@ -108,8 +108,9 @@ def constant(
     out_type: TableType | TensorType
 
     if isinstance(data, TableLike):
-        data_bytes = table_utils.dataframe_to_csv(data)
-        data_format = "bytes[csv]"
+        format = "parquet"
+        data_bytes = table_utils.encode_table(data, format=format)
+        data_format = f"bytes[{format}]"
         out_type = TableType.from_tablelike(data)
     elif isinstance(data, ScalarType):
         out_type = TensorType.from_obj(data)