mplang-nightly 0.1.dev269__py3-none-any.whl → 0.1.dev271__py3-none-any.whl

mplang/v1/ops/basic.py

@@ -1,294 +0,0 @@
-# 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 jax.tree_util import PyTreeDef, tree_flatten
-
-from mplang.v1.core import (
-    UINT8,
-    UINT64,
-    MPObject,
-    PFunction,
-    ScalarType,
-    Shape,
-    TableLike,
-    TableType,
-    TensorLike,
-    TensorType,
-)
-from mplang.v1.ops.base import stateless_mod
-from mplang.v1.utils import table_utils
-
-_BASIC_MOD = stateless_mod("basic")
-
-
-@_BASIC_MOD.simple_op()
-def identity(x: TensorType) -> TensorType:
-    """Return the input type unchanged.
-
-    Args:
-        x: The input tensor type. If called with an MPObject, the value is
-            captured positionally; the kernel sees only the type.
-
-    Returns:
-        The same type as ``x``.
-    """
-    return x
-
-
-@_BASIC_MOD.simple_op()
-def read(*, path: str, ty: TensorType) -> TensorType:
-    """Declare reading a value of type ``ty`` from ``path`` (type-only).
-
-    Args:
-        path: Non-empty path or URI to read from (stored as an attribute).
-        ty: The expected output type/schema.
-
-    Returns:
-        Exactly ``ty``.
-
-    Raises:
-        ValueError: If ``path`` is empty.
-        TypeError: If ``ty`` is not a TensorType or TableType.
-    """
-    if not isinstance(path, str) or path == "":
-        raise ValueError("path must be a non-empty string")
-    if not isinstance(ty, (TensorType, TableType)):
-        raise TypeError("ty must be a TensorType or TableType")
-    # typed_op will attach 'path' as an attribute and build the PFunction
-    return ty
-
-
-@_BASIC_MOD.simple_op()
-def write(x: TensorType, *, path: str) -> TensorType:
-    """Declare writing the input value to ``path`` and return the same type.
-
-    Args:
-        x: The value's type to be written; values are captured positionally.
-        path: Destination path or URI (attribute).
-
-    Returns:
-        The same type as ``x``.
-    """
-    return x
-
-
-@_BASIC_MOD.op_def()
-def constant(
-    data: TensorLike | ScalarType | TableLike,
-) -> tuple[PFunction, list[MPObject], PyTreeDef]:
-    """Embed a literal tensor/table and return the full triad.
-
-    Args:
-        data: Constant payload. Supports scalars, array-like tensors, or
-            table-like dataframes.
-
-    Returns:
-        Tuple[PFunction, list[MPObject], PyTreeDef]:
-        - PFunction: ``fn_type='basic.constant'`` with one output whose type
-            matches ``data``; payload serialized via ``data_bytes`` with
-            ``data_format`` ('bytes[numpy]' or 'bytes[csv]').
-        - list[MPObject]: Empty (no inputs captured).
-        - PyTreeDef: Output tree (single leaf).
-    """
-    import numpy as np
-
-    data_bytes: bytes
-    out_type: TableType | TensorType
-
-    if isinstance(data, TableLike):
-        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)
-        np_data = np.array(data)
-        data_bytes = np_data.tobytes()
-        data_format = "bytes[numpy]"
-    else:
-        if hasattr(data, "tobytes"):
-            out_type = TensorType.from_obj(data)
-            data_bytes = data.tobytes()  # type: ignore[attr-defined]
-        else:
-            np_data = np.array(data)
-            out_type = TensorType.from_obj(np_data)
-            data_bytes = np_data.tobytes()
-        data_format = "bytes[numpy]"
-
-    pfunc = PFunction(
-        fn_type="basic.constant",
-        ins_info=(),
-        outs_info=(out_type,),
-        data_bytes=data_bytes,
-        data_format=data_format,
-    )
-    _, treedef = tree_flatten(out_type)
-    return pfunc, [], treedef
-
-
-@_BASIC_MOD.simple_op()
-def rank() -> TensorType:
-    """Return the scalar UINT64 tensor type for the current party rank.
-
-    Returns:
-        A scalar ``UINT64`` tensor type (shape ``()``).
-    """
-    return TensorType(UINT64, ())
-
-
-@_BASIC_MOD.simple_op()
-def prand(*, shape: Shape = ()) -> TensorType:
-    """Declare a private random UINT64 tensor with the given shape.
-
-    Args:
-        shape: Output tensor shape. Defaults to ``()``.
-
-    Returns:
-        A ``UINT64`` tensor type with the specified shape.
-    """
-    return TensorType(UINT64, shape)
-
-
-@_BASIC_MOD.simple_op()
-def debug_print(
-    x: TensorType | TableType, *, prefix: str = ""
-) -> TableType | TensorType:
-    """Print a value at runtime and return the same type.
-
-    Args:
-        x: The value to print (captured positionally; kernel sees only type).
-        prefix: Optional text prefix for the printed output.
-
-    Returns:
-        The same type as ``x``.
-    """
-    return x
-
-
-@_BASIC_MOD.simple_op()
-def pack(x: TensorType | TableType) -> TensorType:
-    """Serialize a tensor/table into a byte vector (type-only).
-
-    Args:
-        x: Input type to pack.
-
-    Returns:
-        A ``UINT8`` tensor type with shape ``(-1,)`` (length decided at runtime).
-
-    Raises:
-        TypeError: If ``x`` is not a TensorType or TableType.
-    """
-
-    if not isinstance(x, (TensorType, TableType)):
-        raise TypeError("pack expects TensorType or TableType input")
-
-    return TensorType(UINT8, (-1,))
-
-
-@_BASIC_MOD.simple_op()
-def unpack(b: TensorType, *, out_ty: TensorType | TableType) -> TensorType | TableType:
-    """Deserialize a byte vector into the explicit output type.
-
-    Args:
-        b: Byte vector type. Must be ``UINT8`` with shape ``(N,)`` (``N`` may be
-            ``-1``).
-        out_ty: Resulting type/schema after unpacking.
-
-    Returns:
-        Exactly ``out_ty``.
-
-    Raises:
-        TypeError: If ``out_ty`` is not a TensorType/TableType, or if ``b`` is
-            not a 1-D UINT8 tensor.
-    """
-
-    if not isinstance(out_ty, (TensorType, TableType)):
-        raise TypeError("out_ty must be TensorType or TableType")
-
-    if b.dtype != UINT8 or len(b.shape) != 1:
-        raise TypeError("unpack expects a 1-D UINT8 tensor")
-
-    return out_ty
-
-
-@_BASIC_MOD.simple_op()
-def table_to_tensor(table: TableType, *, number_rows: int) -> TensorType:
-    """Convert a homogeneous-typed table to a dense 2D tensor.
-
-    Args:
-        table: Input table whose columns all share the same dtype.
-        number_rows: Number of rows in the resulting tensor. Must be ``>= 0``.
-
-    Returns:
-        A rank-2 tensor with dtype equal to the table column dtype and shape
-        ``(number_rows, table.num_columns())``.
-
-    Raises:
-        ValueError: If the table is empty or ``number_rows < 0``.
-        TypeError: If the table has heterogeneous column dtypes or ``number_rows``
-            is not an int.
-    """
-    if table.num_columns() == 0:
-        raise ValueError("Cannot pack empty table")
-    col_dtypes = list(table.column_types())
-    first = col_dtypes[0]
-    if not all(dt == first for dt in col_dtypes[1:]):
-        raise TypeError(
-            "Heterogeneous dtypes; perform casting upstream before table_to_tensor"
-        )
-    if not isinstance(number_rows, int):
-        raise TypeError("number_rows must be an int")
-    if number_rows < 0:
-        raise ValueError("number_rows must be >= 0")
-    shape = (number_rows, table.num_columns())
-    return TensorType(first, shape)  # type: ignore[arg-type]
-
-
-@_BASIC_MOD.simple_op()
-def tensor_to_table(tensor: TensorType, *, column_names: list[str]) -> TableType:
-    """Convert a rank-2 tensor into a table with named columns.
-
-    Args:
-        tensor: Rank-2 tensor with shape ``(N, F)``.
-        column_names: List of unique, non-whitespace column names of length ``F``.
-
-    Returns:
-        A table with ``F`` columns named as provided, each with dtype
-        ``tensor.dtype``.
-
-    Raises:
-        TypeError: If ``tensor`` is not rank-2, or if any column name is not a
-            string.
-        ValueError: If names are empty/whitespace, duplicated, or length != ``F``.
-    """
-    if len(tensor.shape) != 2:
-        raise TypeError("tensor_to_table expects a rank-2 tensor (N,F)")
-    n_cols = tensor.shape[1]
-    if not column_names:
-        raise ValueError("column_names required (non-empty)")
-    if len(column_names) != n_cols:
-        raise ValueError("column_names length must match tensor second dim")
-    for i, name in enumerate(column_names):
-        if not isinstance(name, str):
-            raise TypeError(f"column_names[{i}] must be str, got {type(name).__name__}")
-        if name == "" or name.strip() == "":
-            raise ValueError("column names must be non-empty and not whitespace-only")
-    seen: set[str] = set()
-    for name in column_names:
-        if name in seen:
-            raise ValueError(f"duplicate column name: {name!r}")
-        seen.add(name)
-    col_types = [tensor.dtype] * n_cols
-    return TableType.from_pairs(list(zip(column_names, col_types, strict=True)))

mplang/v1/ops/crypto.py

@@ -1,262 +0,0 @@
-# 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.
-
-"""
-Crypto frontend operations: operation signatures, types, and high-level semantics.
-
-Scope and contracts:
-- This module defines portable API shapes; it does not implement cryptography.
-- Backends execute the operations and must meet the security semantics required
-    by the deployment (confidentiality, authenticity, correctness, etc.).
-- The enc/dec API in this frontend uses a conventional 12-byte nonce prefix
-    (ciphertext = nonce || payload), and dec expects that format. Other security
-    properties (e.g., AEAD) are backend responsibilities.
-"""
-
-from __future__ import annotations
-
-from jax.tree_util import PyTreeDef, tree_flatten
-
-from mplang.v1.core import UINT8, TensorType
-from mplang.v1.core.mpobject import MPObject
-from mplang.v1.core.pfunc import PFunction
-from mplang.v1.ops.base import stateless_mod
-
-_CRYPTO_MOD = stateless_mod("crypto")
-
-
-def _get_algo_overhead(algo: str) -> int:
-    """Get ciphertext overhead for a given encryption algorithm.
-
-    Args:
-        algo: Encryption algorithm identifier
-
-    Returns:
-        int: Number of overhead bytes added to plaintext length
-    """
-    overhead_map = {
-        "aes-ctr": 16,  # nonce only (legacy compatibility)
-        "aes-gcm": 28,  # nonce(12) + tag(16) for AES-GCM
-        "sm4-gcm": 28,  # nonce(12) + tag(16) for SM4-GCM
-    }
-
-    if algo not in overhead_map:
-        # return unknown overhead as -1
-        return -1
-    return overhead_map[algo]
-
-
-@_CRYPTO_MOD.simple_op()
-def keygen(*, length: int = 32) -> TensorType:
-    """Generate random bytes for symmetric keys or generic randomness.
-
-    API: keygen(length: int = 32) -> key: u8[length]
-
-    Notes:
-    - Frontend defines the type/shape; backend provides randomness.
-    - Raises ValueError when length <= 0.
-    """
-    if length <= 0:
-        raise ValueError("length must be > 0")
-    return TensorType(UINT8, (length,))
-
-
-@_CRYPTO_MOD.op_def()
-def enc(
-    plaintext: MPObject, key: MPObject, algo: str = "aes-ctr"
-) -> tuple[PFunction, list[MPObject], PyTreeDef]:
-    """Symmetric encryption with algorithm-aware output sizing.
-
-    API: enc(plaintext: u8[N], key: u8[M], *, algo: str = "aes-ctr") -> ciphertext: u8[N + overhead]
-
-    Supported algorithms and overhead:
-    - "aes-ctr": 16 bytes (nonce only, legacy compatibility)
-    - "aes-gcm": 28 bytes (nonce + 16-byte authentication tag)
-    - "sm4-gcm": 28 bytes (nonce + 16-byte authentication tag)
-
-    The algo parameter is stored in the PFunction attributes for backend use.
-    """
-    pt_ty = plaintext
-    if pt_ty.dtype != UINT8:
-        raise TypeError("enc expects UINT8 plaintext")
-    if len(pt_ty.shape) != 1:
-        raise TypeError("enc expects 1-D plaintext")
-
-    # Validate and get overhead for the specified algorithm
-    overhead = _get_algo_overhead(algo)
-    length = pt_ty.shape[0]
-    if length >= 0 and overhead >= 0:
-        outs_info = (TensorType(UINT8, (length + overhead,)),)
-    else:
-        # Unknown length or overhead, return dynamic length
-        outs_info = (TensorType(UINT8, (-1,)),)
-
-    ins_info = (TensorType.from_obj(pt_ty), TensorType.from_obj(key))
-    pfunc = PFunction(
-        fn_type="crypto.enc",
-        ins_info=ins_info,
-        outs_info=outs_info,
-        algo=algo,
-    )
-    _, treedef = tree_flatten(outs_info[0])
-    return pfunc, [plaintext, key], treedef
-
-
-@_CRYPTO_MOD.op_def()
-def dec(
-    ciphertext: MPObject, key: MPObject, algo: str = "aes-ctr"
-) -> tuple[PFunction, list[MPObject], PyTreeDef]:
-    """Symmetric decryption with algorithm-aware input sizing.
-
-    API: dec(ciphertext: u8[N + overhead], key: u8[M], *, algo: str = "aes-ctr") -> plaintext: u8[N]
-
-    Supported algorithms and overhead:
-    - "aes-ctr": 16 bytes (nonce only, legacy compatibility)
-    - "aes-gcm": 28 bytes (nonce + 16-byte authentication tag)
-    - "sm4-gcm": 28 bytes (nonce + 16-byte authentication tag)
-
-    The algo parameter is stored in the PFunction attributes for backend use.
-    Backend is responsible for parsing the ciphertext format according to algo.
-    """
-    ct_ty = ciphertext
-    if ct_ty.dtype != UINT8:
-        raise TypeError("dec expects UINT8 ciphertext")
-    if len(ct_ty.shape) != 1:
-        raise TypeError("dec expects 1-D ciphertext")
-
-    # Validate and get overhead for the specified algorithm
-    overhead = _get_algo_overhead(algo)
-    length = ct_ty.shape[0]
-
-    # Validate minimum ciphertext length
-    if length >= 0 and overhead >= 0 and length < overhead:
-        raise TypeError(
-            f"dec expects ciphertext with at least {overhead} bytes for algo='{algo}', but got {length} bytes"
-        )
-
-    # Compute output plaintext length
-    if length >= 0 and overhead >= 0:
-        outs_info = (TensorType(UINT8, (length - overhead,)),)
-    else:
-        # Unknown length or overhead, return dynamic length
-        outs_info = (TensorType(UINT8, (-1,)),)
-
-    ins_info = (TensorType.from_obj(ct_ty), TensorType.from_obj(key))
-    pfunc = PFunction(
-        fn_type="crypto.dec",
-        ins_info=ins_info,
-        outs_info=outs_info,
-        algo=algo,
-    )
-    _, treedef = tree_flatten(outs_info[0])
-    return pfunc, [ciphertext, key], treedef
-
-
-@_CRYPTO_MOD.op_def()
-def kem_keygen(suite: str = "x25519") -> tuple[PFunction, list[MPObject], PyTreeDef]:
-    """KEM-style keypair generation: returns (sk, pk) bytes.
-
-    API: kem_keygen(suite: str = "x25519") -> (sk: u8[32], pk: u8[32])
-
-    The suite parameter is stored in the PFunction attributes for backend use.
-    """
-    if suite == "x25519":
-        sk_ty = TensorType(UINT8, (32,))
-        pk_ty = TensorType(UINT8, (32,))
-    else:
-        # Unknown suite, return dynamic lengths
-        sk_ty = TensorType(UINT8, (-1,))
-        pk_ty = TensorType(UINT8, (-1,))
-    outs_info = (sk_ty, pk_ty)
-
-    pfunc = PFunction(
-        fn_type="crypto.kem_keygen",
-        ins_info=(),
-        outs_info=outs_info,
-        suite=suite,
-    )
-    _, treedef = tree_flatten(outs_info)
-    return pfunc, [], treedef
-
-
-@_CRYPTO_MOD.op_def()
-def kem_derive(
-    sk: MPObject, peer_pk: MPObject, suite: str = "x25519"
-) -> tuple[PFunction, list[MPObject], PyTreeDef]:
-    """KEM-style shared secret derivation: returns secret bytes.
-
-    API: kem_derive(sk: u8[32], peer_pk: u8[32], suite: str = "x25519") -> secret: u8[32]
-
-    The suite parameter is stored in the PFunction attributes for backend use.
-    """
-    # Validate input types
-    if sk.dtype != UINT8:
-        raise TypeError("kem_derive expects UINT8 secret key")
-    if peer_pk.dtype != UINT8:
-        raise TypeError("kem_derive expects UINT8 peer public key")
-    if len(sk.shape) != 1 or len(peer_pk.shape) != 1:
-        raise TypeError("kem_derive expects 1-D inputs")
-
-    if suite == "x25519":
-        if sk.shape[0] != 32 or peer_pk.shape[0] != 32:
-            raise TypeError("kem_derive expects 32-byte keys for suite 'x25519'")
-        secret_ty = TensorType(UINT8, (32,))
-    else:
-        # Unknown suite, return dynamic length
-        secret_ty = TensorType(UINT8, (-1,))
-    outs_info = (secret_ty,)
-
-    ins_info = (TensorType.from_obj(sk), TensorType.from_obj(peer_pk))
-    pfunc = PFunction(
-        fn_type="crypto.kem_derive",
-        ins_info=ins_info,
-        outs_info=outs_info,
-        suite=suite,
-    )
-    _, treedef = tree_flatten(outs_info[0])
-    return pfunc, [sk, peer_pk], treedef
-
-
-@_CRYPTO_MOD.op_def()
-def hkdf(
-    secret: MPObject, info: str, hash: str = "SHA-256"
-) -> tuple[PFunction, list[MPObject], PyTreeDef]:
-    """HKDF-style key derivation: returns a 32-byte key.
-
-    API: hkdf(secret: u8[N], info: str, hash: str = "SHA-256") -> key: u8[32]
-
-    The hash parameter is stored in the PFunction attributes for backend use.
-    """
-    # Validate input types
-    if secret.dtype != UINT8:
-        raise TypeError("hkdf expects UINT8 secret")
-    if len(secret.shape) != 1:
-        raise TypeError("hkdf expects 1-D secret")
-
-    if hash == "SHA-256" or hash == "SM3":
-        outs_info = (TensorType(UINT8, (32,)),)
-    else:
-        # Unknown hash, return dynamic length
-        outs_info = (TensorType(UINT8, (-1,)),)
-
-    ins_info = (TensorType.from_obj(secret),)
-    pfunc = PFunction(
-        fn_type="crypto.hkdf",
-        ins_info=ins_info,
-        outs_info=outs_info,
-        hash=hash,
-        info=info,
-    )
-    _, treedef = tree_flatten(outs_info[0])
-    return pfunc, [secret], treedef