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

mplang/v1/core/mpobject.py

@@ -1,117 +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 __future__ import annotations
-
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
-
-from mplang.v1.core.dtypes import DType
-from mplang.v1.core.mask import Mask
-from mplang.v1.core.mptype import MPType
-from mplang.v1.core.table import TableType
-from mplang.v1.core.tensor import Shape
-
-if TYPE_CHECKING:
-    from mplang.v1.core.cluster import ClusterSpec
-
-
-class MPContext:
-    """The context of an MPObject.
-
-    MPContext is the abstract base class for all execution contexts.
-    It only holds the immutable cluster_spec plus lightweight parent/root
-    helpers used to support stack-scoped extension state (attached lazily by
-    external features on the root context).
-    """
-
-    def __init__(self, cluster_spec: ClusterSpec, *, parent: MPContext | None = None):
-        if cluster_spec is None:
-            raise ValueError("cluster_spec cannot be None")
-        self.cluster_spec = cluster_spec
-        # Parent link enables stack-scoped state sharing: ephemeral child contexts
-        # (e.g. short-lived tracing) can delegate to a stable root without relying
-        # on process-wide globals.
-        self._parent: MPContext | None = parent
-
-    # Basic topology helpers
-    def world_size(self) -> int:
-        return len(self.cluster_spec.nodes)
-
-    @property
-    def parent(self) -> MPContext | None:
-        """Direct parent context or None if this is root."""
-        return self._parent
-
-    def root(self) -> MPContext:
-        """Return the root context (follow parent chain)."""
-        ctx: MPContext = self
-        visited: set[int] = set()
-        while ctx._parent is not None:
-            if id(ctx) in visited:
-                raise RuntimeError("Cycle detected in MPContext parent chain")
-            visited.add(id(ctx))
-            ctx = ctx._parent
-        return ctx
-
-
-class MPObject(ABC):
-    """The base class for all objects in mp-system."""
-
-    @property
-    @abstractmethod
-    def mptype(self) -> MPType:
-        """The type information of the object.
-
-        This property is readonly (mandatory) and will be used for JAX compilation
-        to determine the appropriate data type during trace and compilation phases.
-        MPType can be passed between different MPObjects as a value.
-        """
-
-    @property
-    def dtype(self) -> DType:
-        return self.mptype.dtype
-
-    @property
-    def shape(self) -> Shape:
-        return self.mptype.shape
-
-    @property
-    def schema(self) -> TableType:
-        """The table schema of the object.
-
-        Only available for table types.
-        """
-        return self.mptype.schema
-
-    @property
-    def pmask(self) -> Mask | None:
-        return self.mptype.pmask
-
-    @property
-    def attrs(self) -> dict[str, Any]:
-        return self.mptype.attrs
-
-    @property
-    @abstractmethod
-    def ctx(self) -> MPContext:
-        """Return the context of the object."""
-
-
-# Forward docstrings from MPType to MPObject
-MPObject.dtype.__doc__ = MPType.dtype.__doc__
-MPObject.shape.__doc__ = MPType.shape.__doc__
-MPObject.schema.__doc__ = MPType.schema.__doc__
-MPObject.pmask.__doc__ = MPType.pmask.__doc__
-MPObject.attrs.__doc__ = MPType.attrs.__doc__

mplang/v1/core/mptype.py

@@ -1,407 +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 __future__ import annotations
-
-import copy
-from typing import TYPE_CHECKING, Any
-
-import numpy as np
-
-if TYPE_CHECKING:
-    from mplang.v1.core.mpobject import MPObject
-
-from mplang.v1.core.dtypes import STRING, DType
-from mplang.v1.core.mask import Mask
-from mplang.v1.core.table import TableLike, TableType
-from mplang.v1.core.tensor import ScalarType, Shape, TensorLike, TensorType
-
-# basic type aliases
-Rank = int
-
-
-class MPType:
-    """A type that describes the type information of an MPObject."""
-
-    _type: TensorType | TableType
-    _pmask: Mask | None
-    _attrs: dict[str, Any]
-
-    def __init__(
-        self,
-        type_info: TensorType | TableType,
-        pmask: Mask | None = None,
-        attrs: dict[str, Any] | None = None,
-    ):
-        """Initialize MPType.
-
-        Args:
-            type_info: The type information (TensorType for tensors, TableType for tables).
-            pmask: The party mask, used for compile/trace time determine which party holds the object.
-            attrs: Attributes are key-value pairs that can be used to store additional information about the object.
-        """
-        self._type = type_info
-        self._pmask = pmask
-        # Ensure attrs is a copy
-        self._attrs = copy.copy(attrs) if attrs is not None else {}
-
-    @classmethod
-    def tensor(
-        cls,
-        dtype: DType | Any,
-        shape: Shape,
-        pmask: int | Mask | None = None,
-        **attrs: Any,
-    ) -> MPType:
-        """Create a tensor type.
-
-        Args:
-            dtype: The data type of the tensor.
-            shape: The shape of the tensor.
-            pmask: The party mask.
-            **attrs: Additional attributes.
-
-        Returns:
-            MPType instance for tensor.
-
-        Raises:
-            ValueError: If dtype is table-only.
-        """
-        # Convert dtype to DType if needed and validate
-        if not isinstance(dtype, DType):
-            dtype = DType.from_any(dtype)
-
-        # Ensure tensor types don't use table-only dtypes
-        if dtype.is_table_only:
-            raise ValueError(
-                f"Data type '{dtype.name}' is only supported in tables, "
-                f"not in tensors. Use table types for string, date, and other "
-                f"non-numeric data types."
-            )
-
-        if isinstance(pmask, int):
-            pmask = Mask.from_int(pmask)
-
-        tensor_info = TensorType(dtype, shape)
-        return cls(tensor_info, pmask, attrs)
-
-    @classmethod
-    def table(
-        cls,
-        schema: TableType | dict[str, DType],
-        pmask: int | Mask | None = None,
-        **attrs: Any,
-    ) -> MPType:
-        """Create a table type.
-
-        Args:
-            schema: The table schema or dict mapping column names to types.
-            pmask: The party mask.
-            **attrs: Additional attributes.
-
-        Returns:
-            MPType instance for table.
-        """
-        if isinstance(schema, dict):
-            schema = TableType.from_dict(schema)
-
-        if isinstance(pmask, int):
-            pmask = Mask.from_int(pmask)
-
-        return cls(schema, pmask, attrs)
-
-    @property
-    def is_tensor(self) -> bool:
-        """Check if this is a tensor type."""
-        return isinstance(self._type, TensorType)
-
-    @property
-    def is_table(self) -> bool:
-        """Check if this is a table type."""
-        return isinstance(self._type, TableType)
-
-    @property
-    def dtype(self) -> DType:
-        """The data type of the object.
-
-        This property is readonly (mandatory) and will be used for JAX compilation
-        to determine the appropriate data type during trace and compilation phases.
-
-        Only available for tensor types.
-        """
-        if not isinstance(self._type, TensorType):
-            raise AttributeError("dtype is only available for tensor types")
-        return self._type.dtype
-
-    @property
-    def shape(self) -> Shape:
-        """The shape of the object, represented as a tuple of integers.
-
-        For example, a 2D tensor with shape (3, 4) would be represented as (3, 4).
-        The shape can be empty, which indicates a scalar.
-
-        This property is readonly (mandatory) and will be used for JAX compilation
-        to determine tensor shapes during trace and compilation phases.
-
-        Only available for tensor types.
-        """
-        if not isinstance(self._type, TensorType):
-            raise AttributeError("shape is only available for tensor types")
-        return self._type.shape
-
-    @property
-    def schema(self) -> TableType:
-        """The table schema.
-
-        Only available for table types.
-        """
-        if not isinstance(self._type, TableType):
-            raise AttributeError("schema is only available for table types")
-        return self._type
-
-    @property
-    def pmask(self) -> Mask | None:
-        """The party mask indicating which parties hold the data.
-
-        Value interpretation:
-        - When not None: A bitmask where the i'th bit is 1 if the i'th party holds
-          the data, and 0 otherwise. For example, 0b1101 means parties 0, 2, and 3
-          hold the data, while party 1 does not.
-        - When None: Party ownership is unknown at compile/trace time and will be
-          completely determined at runtime.
-
-        Semantic meaning:
-        This mask can be either manually set or deduced by primitive functions during
-        compilation/tracing. When None, it does NOT imply either a full mask (all
-        parties) or zero mask (no parties) - the actual ownership pattern is entirely
-        runtime-dependent.
-        """
-        return self._pmask
-
-    @property
-    def attrs(self) -> dict[str, Any]:
-        """Attributes are key-value pairs that can be used to store additional
-        information about the object."""
-        return self._attrs
-
-    def raw_type(self) -> TensorType | TableType:
-        """Get the raw type information (TensorType or TableType)."""
-        return self._type
-
-    def set_attr(self, key: str, value: Any) -> None:
-        """Set an attribute for this type."""
-        self._attrs[key] = value
-
-    def get_attr(self, key: str, default: Any = None) -> Any:
-        """Get an attribute for this type."""
-        return self._attrs.get(key, default)
-
-    def __repr__(self) -> str:
-        """String representation of MPType.
-
-        Schema:
-        - For tensor: dtype[shape]<pmask>{other_attrs}
-        - For table: Tbl(col1:type1, col2:type2)<pmask>{other_attrs}
-
-        Examples:
-        - u64                        # scalar uint64
-        - f32[3, 2]                 # 3x2 float32 tensor
-        - f16[3]<3>                 # float16 vector with pmask=3
-        - u32[5, 5]<F>{device="P0"} # uint32 matrix with pmask=15 and device attr
-        - Tbl(id:i64, name:str)      # table with id and name columns
-        """
-        if isinstance(self._type, TensorType):
-            # Start with short dtype name
-            ret = self._type.dtype.short_name()
-
-            # Add shape if not scalar
-            if self._type.shape:
-                shape_str = ", ".join(str(d) for d in self._type.shape)
-                ret += f"[{shape_str}]"
-        else:  # TableType
-            cols = ", ".join(
-                f"{name}:{dtype.short_name()}" for name, dtype in self._type.columns
-            )
-            ret = f"Tbl({cols})"
-
-        # Add pmask in angle brackets if present
-        if self._pmask is not None:
-            ret += f"<{self._pmask:X}>"
-
-        # Add other attributes in curly braces if any
-        if self._attrs:
-            attrs_list = []
-            for key, value in self._attrs.items():
-                if isinstance(value, str):
-                    attrs_list.append(f'{key}="{value}"')
-                else:
-                    attrs_list.append(f"{key}={value}")
-            ret += "{" + ", ".join(attrs_list) + "}"
-
-        return ret
-
-    def __eq__(self, other: object) -> bool:
-        """Check if two MPType objects are equal."""
-        if not isinstance(other, MPType):
-            return False
-        return (
-            self._type == other._type and self._pmask == other._pmask
-            # and self._attrs == other._attrs # TODO(jint): attrs should be optional
-        )
-
-    def __hash__(self) -> int:
-        """Compute hash for MPType objects."""
-        # Make attrs hashable by converting to frozenset of items
-        attrs_hash = hash(frozenset(self._attrs.items())) if self._attrs else 0
-        return hash((
-            self._type,
-            self._pmask,
-            attrs_hash,
-        ))
-
-    def isInstance(self, obj: MPObject) -> bool:
-        """Check if the given object is an instance of this MPType."""
-        # Import here to avoid circular import
-        from mplang.v1.core.mpobject import MPObject
-
-        if not isinstance(obj, MPObject):
-            return False
-
-        # Check if the object's type matches this type
-        obj_type = obj.mptype
-        if type(self._type) is not type(obj_type._type):
-            return False
-
-        if self._type != obj_type._type:
-            return False
-
-        # Check attributes
-        if self._attrs:
-            if not isinstance(obj.attrs, dict):
-                return False
-            for k, v in self._attrs.items():
-                if k not in obj.attrs or obj.attrs[k] != v:
-                    return False
-        return True
-
-    def to_numpy(self) -> np.dtype:
-        """Convert to NumPy dtype for compatibility.
-
-        Only available for tensor types.
-        """
-        if not isinstance(self._type, TensorType):
-            raise AttributeError("to_numpy is only available for tensor types")
-        return self._type.to_numpy()
-
-    @staticmethod
-    def _create_tensor_info(obj: TensorLike | ScalarType) -> TensorType:
-        """Helper method to create TensorType from tensor-like objects."""
-        if isinstance(obj, ScalarType):
-            return TensorType(DType.from_python_type(type(obj)), ())
-        elif isinstance(obj, TensorLike):
-            return TensorType(DType.from_any(obj.dtype), obj.shape)
-        elif isinstance(obj, list | tuple):
-            # Convert lists/tuples to numpy arrays for compatibility
-            arr = np.array(obj)
-            return TensorType(DType.from_any(arr.dtype), arr.shape)
-        else:
-            raise TypeError(f"Unsupported type: {type(obj)}.")
-
-    @classmethod
-    def from_tensor(
-        cls,
-        obj: TensorLike | ScalarType,
-        pmask: Mask | None = None,
-        **kwargs: Any,
-    ) -> MPType:
-        """Create MPType from tensor-like object.
-
-        Args:
-            obj: Tensor-like object or scalar.
-            pmask: The party mask.
-            **kwargs: Additional attributes.
-
-        Returns:
-            MPType instance for tensor.
-        """
-        attrs = copy.copy(kwargs)
-        tensor_info = cls._create_tensor_info(obj)
-        return cls(tensor_info, pmask, attrs)
-
-    @classmethod
-    def from_mpobj(cls, obj: MPObject) -> MPType:
-        """Create MPType from MPObject.
-
-        Args:
-            obj: MPObject instance.
-
-        Returns:
-            MPType instance with same type as the object.
-        """
-        # assume obj is MPObject-like
-        obj_type = obj.mptype
-        return cls(obj_type._type, obj.pmask, copy.copy(obj.attrs))
-
-    @classmethod
-    def from_obj(cls, obj: Any, pmask: Mask | None = None, **attrs: Any) -> MPType:
-        """Create MPType from any object, automatically inferring the type.
-
-        Args:
-            obj: Object to create type from.
-            pmask: The party mask.
-            **attrs: Additional attributes.
-
-        Returns:
-            MPType instance.
-
-        Raises:
-            TypeError: If object type cannot be inferred.
-            NotImplementedError: For table objects (not yet implemented).
-        """
-        # Check if it's a table-like object using the TableLike protocol
-        if isinstance(obj, TableLike):
-            # For TableLike objects, try to extract schema information
-            try:
-                import pandas as pd
-
-                if isinstance(obj, pd.DataFrame):
-                    from mplang.v1.core.dtypes import DType
-
-                    schema_dict = {}
-                    for col_name in obj.columns:
-                        pandas_dtype = obj[col_name].dtype
-                        # Convert pandas dtype to DType
-                        if pandas_dtype.kind in (
-                            "O",
-                            "U",
-                            "S",
-                        ):  # object, unicode, string
-                            schema_dict[col_name] = (
-                                DType.from_numpy(pandas_dtype)
-                                if pandas_dtype.kind != "O"
-                                else STRING
-                            )
-                        else:
-                            schema_dict[col_name] = DType.from_numpy(pandas_dtype)
-                    schema = TableType.from_dict(schema_dict)
-                    return cls(schema, pmask, attrs)
-            except ImportError:
-                pass
-            # For other table-like objects without pandas
-            raise NotImplementedError(
-                "Table object detection for non-pandas objects not fully implemented yet"
-            )
-
-        # Otherwise treat as tensor-like
-        return cls.from_tensor(obj, pmask, **attrs)

mplang/v1/core/pfunc.py

@@ -1,130 +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 __future__ import annotations
-
-import copy
-from collections.abc import Sequence
-from types import MappingProxyType
-from typing import Any
-
-from mplang.v1.core.table import TableType
-from mplang.v1.core.tensor import TensorType
-
-__all__ = [
-    "PFunction",
-    "get_fn_name",
-]
-
-
-class PFunction:
-    """A Party Function represents a computation unit that can be executed by a single party.
-
-    PFunction serves as a unified interface for describing single-party computations
-    in multi-party computing scenarios. It can represent both:
-    1. Built-in operations (e.g., "spu.makeshares", "basic.read")
-    2. User-defined programmable functions with custom code
-
-    The PFunction accepts a list of typed inputs (TensorType/TableType). For
-    backend-only handles (e.g., crypto keys), use a sentinel TensorType
-    of UINT8 with shape (-1, 0) to indicate the argument should bypass
-    structural validation at runtime. Outputs should likewise use concrete
-    TensorType/TableType specs. PFunction can be:
-    - Expressed and defined in the mplang frontend
-    - Serialized for transmission between components
-    - Interpreted and executed by backend runtime engines
-
-    Args:
-        fn_type: The type/category identifier of this PFunction, indicating which
-            backend or handler should process it (e.g., "spu.makeshares", "basic.read",
-            "mlir.stablehlo"). This serves as a routing mechanism for execution.
-        ins_info: Type information for input parameters (TensorType or TableType)
-        outs_info: Type information for output values (TensorType or TableType)
-        fn_name: Optional name of the function. For programmable functions, this is
-            the user-defined function name. For built-in operations, this may be
-            None or a descriptive identifier.
-        fn_text: Optional serialized function body. For programmable functions, this
-            contains the actual code (e.g., MLIR, bytecode, source code). For built-in
-            operations, this is typically None.
-        **kwargs: Additional attributes and metadata specific to the function type.
-            These are used to pass execution parameters, configuration, and context
-            information to the backend handlers.
-    """
-
-    # Required fields - these define the core execution context
-    fn_type: str  # Unique identifier for backend routing
-    ins_info: tuple[TensorType | TableType, ...]
-    outs_info: tuple[TensorType | TableType, ...]
-
-    # Optional fields for programmable functions
-    fn_name: str | None  # Function name (for programmable functions)
-    fn_text: str | None  # Function body/code (for programmable functions)
-
-    # Custom attributes and metadata
-    attrs: MappingProxyType[str, Any]  # Execution parameters and metadata
-
-    def __init__(
-        self,
-        fn_type: str,
-        ins_info: Sequence[TensorType | TableType],
-        outs_info: Sequence[TensorType | TableType],
-        *,
-        fn_name: str | None = None,
-        fn_text: str | None = None,
-        **kwargs: Any,
-    ):
-        self.fn_type = fn_type
-        self.fn_name = fn_name
-        self.fn_text = fn_text
-        self.ins_info = tuple(ins_info)
-        self.outs_info = tuple(outs_info)
-        # Make attrs immutable to ensure PFunction immutability
-        # Create a copy first, then wrap it in MappingProxyType
-        self.attrs = MappingProxyType(copy.copy(kwargs))
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}({self.fn_type}, {self.fn_name})"
-
-    def __hash__(self) -> int:
-        return hash((
-            self.fn_type,
-            self.fn_name,
-            self.fn_text,
-            self.ins_info,
-            self.outs_info,
-            frozenset(self.attrs.items()),
-        ))
-
-    def __eq__(self, other: object) -> bool:
-        """Check equality between PFunction instances."""
-        if not isinstance(other, PFunction):
-            return False
-
-        return (
-            self.fn_type == other.fn_type
-            and self.fn_name == other.fn_name
-            and self.fn_text == other.fn_text
-            and self.ins_info == other.ins_info
-            and self.outs_info == other.outs_info
-            and self.attrs == other.attrs
-        )
-
-
-def get_fn_name(fn_like: Any) -> str:
-    if hasattr(fn_like, "__name__"):
-        return fn_like.__name__  # type: ignore[no-any-return]
-    if hasattr(fn_like, "func"):
-        # handle partial functions
-        return get_fn_name(fn_like.func)
-    return "unnamed function"