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

mplang/{core → v1/core}/expr/printer.py

@@ -20,8 +20,8 @@ from __future__ import annotations
 
 from typing import Any
 
-from mplang.core.dtype import DType
-from mplang.core.expr.ast import (
+from mplang.v1.core.dtypes import DType
+from mplang.v1.core.expr.ast import (
     AccessExpr,
     CallExpr,
     CondExpr,
@@ -35,10 +35,10 @@ from mplang.core.expr.ast import (
     VariableExpr,
     WhileExpr,
 )
-from mplang.core.expr.visitor import ExprVisitor
-from mplang.core.mptype import MPType
-from mplang.core.pfunc import PFunction
-from mplang.core.tensor import Shape, TensorType
+from mplang.v1.core.expr.visitor import ExprVisitor
+from mplang.v1.core.mptype import MPType
+from mplang.v1.core.pfunc import PFunction
+from mplang.v1.core.tensor import Shape, TensorType
 
 
 class Printer(ExprVisitor):
@@ -50,11 +50,13 @@ class Printer(ExprVisitor):
         compact_format: bool = True,
         *,
         verbose_peval: bool = False,
+        inline_pcall: bool = True,
     ):
         super().__init__()  # Initialize MemorizedVisitor
         self.indent_size = indent_size
         self.compact_format = compact_format
         self.verbose_peval = verbose_peval
+        self.inline_pcall = inline_pcall
         self._cur_indent = 0
         self._output: list[str] = []
         self._visited: dict[Expr, str] = {}
@@ -92,6 +94,7 @@ class Printer(ExprVisitor):
                 body_printer = Printer(
                     indent_size=self.indent_size,
                     compact_format=self.compact_format,
+                    inline_pcall=self.inline_pcall,
                 )
                 func_def_expr.accept(body_printer)
                 regions_str += f"{indent}{r_name}: "
@@ -161,13 +164,9 @@ class Printer(ExprVisitor):
         arg_names = [self._var_name(arg) for arg in expr.args]
         fn_type = expr.pfunc.fn_type
 
-        # for well known builtin functions
-        if fn_type == "builtin.constant":
+        # for well known basic functions
+        if fn_type == "basic.constant":
             return self._print_const(expr.pfunc, expr.mptypes)
-        elif fn_type == "builtin.rank":
-            return self._do_print("prank", [], mptypes=expr.mptypes)
-        elif fn_type == "builtin.prand":
-            return self._do_print("prand", [], mptypes=expr.mptypes)
 
         attrs = {"fn_type": fn_type}
         if expr.pfunc.fn_name:
@@ -209,12 +208,19 @@ class Printer(ExprVisitor):
 
     def visit_call(self, expr: CallExpr) -> str:
         arg_names = [self._var_name(arg) for arg in expr.args]
-        return self._do_print(
-            "pcall",
-            arg_names,
-            regions={"fn": expr.fn},
-            mptypes=expr.mptypes,
-        )
+        if self.inline_pcall:
+            return self._do_print(
+                expr.name,
+                arg_names,
+                mptypes=expr.mptypes,
+            )
+        else:
+            return self._do_print(
+                "pcall",
+                arg_names,
+                regions={"fn": expr.fn},
+                mptypes=expr.mptypes,
+            )
 
     def visit_while(self, expr: WhileExpr) -> str:
         arg_names = [self._var_name(arg) for arg in expr.args]

mplang/{core → v1/core}/expr/transformer.py

@@ -18,7 +18,7 @@ Expression transformer based on visitor pattern.
 
 from collections.abc import Callable
 
-from mplang.core.expr.ast import (
+from mplang.v1.core.expr.ast import (
     AccessExpr,
     CallExpr,
     CondExpr,
@@ -32,7 +32,7 @@ from mplang.core.expr.ast import (
     VariableExpr,
     WhileExpr,
 )
-from mplang.core.expr.visitor import ExprVisitor
+from mplang.v1.core.expr.visitor import ExprVisitor
 
 
 class ExprTransformer(ExprVisitor):
@@ -79,7 +79,7 @@ class ExprTransformer(ExprVisitor):
     def visit_call(self, expr: CallExpr) -> Expr:
         # Transform child expressions first
         transformed_args = [arg.accept(self) for arg in expr.args]
-        new_expr = CallExpr(expr.fn, transformed_args)
+        new_expr = CallExpr(expr.name, expr.fn, transformed_args)
 
         if "call" in self.trans_rules:
             return self.trans_rules["call"](new_expr)

mplang/{core → v1/core}/expr/utils.py

@@ -18,8 +18,8 @@ Utility functions for expression system.
 
 from collections.abc import Sequence
 
-from mplang.core.mask import Mask
-from mplang.core.mptype import TensorLike
+from mplang.v1.core.mask import Mask
+from mplang.v1.core.mptype import TensorLike
 
 
 def type_equal(*args: TensorLike) -> bool:

mplang/{core → v1/core}/expr/visitor.py

@@ -22,7 +22,7 @@ from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-    from mplang.core.expr.ast import (
+    from mplang.v1.core.expr.ast import (
         AccessExpr,
         CallExpr,
         CondExpr,

mplang/{core → v1/core}/expr/walk.py

@@ -32,7 +32,7 @@ from collections import deque
 from collections.abc import Callable, Iterable, Iterator, Sequence
 from typing import cast
 
-from mplang.core.expr.ast import (
+from mplang.v1.core.expr.ast import (
     AccessExpr,
     CallExpr,
     CondExpr,

mplang/{core → v1/core}/interp.py

@@ -25,12 +25,12 @@ from abc import abstractmethod
 from collections.abc import Sequence
 from typing import Any, cast
 
-from mplang.core.cluster import ClusterSpec
-from mplang.core.expr.ast import Expr, VariableExpr
-from mplang.core.mpobject import MPContext, MPObject
-from mplang.core.mptype import MPType, TensorLike
-from mplang.core.tracer import TracedFunction
-from mplang.utils.func_utils import var_demorph, var_morph
+from mplang.v1.core.cluster import ClusterSpec
+from mplang.v1.core.expr.ast import Expr, VariableExpr
+from mplang.v1.core.mpobject import MPContext, MPObject
+from mplang.v1.core.mptype import MPType, TensorLike
+from mplang.v1.core.tracer import TracedFunction
+from mplang.v1.utils.func_utils import var_demorph, var_morph
 
 
 # TODO(jint): Should we use inheritance or composition here?

mplang/{core → v1/core}/mpir.py

@@ -32,9 +32,9 @@ from typing import Any
 import numpy as np
 import spu.libspu as spu_api
 
-from mplang.core.dtype import DATE, JSON, STRING, TIME, TIMESTAMP, DType
-from mplang.core.expr import Expr, FuncDefExpr
-from mplang.core.expr.ast import (
+from mplang.v1.core.dtypes import DATE, JSON, STRING, TIME, TIMESTAMP, DType
+from mplang.v1.core.expr import Expr, FuncDefExpr
+from mplang.v1.core.expr.ast import (
     AccessExpr,
     CallExpr,
     CondExpr,
@@ -46,13 +46,13 @@ from mplang.core.expr.ast import (
     VariableExpr,
     WhileExpr,
 )
-from mplang.core.expr.walk import walk
-from mplang.core.mask import Mask
-from mplang.core.mptype import MPType
-from mplang.core.pfunc import PFunction
-from mplang.core.table import TableType
-from mplang.core.tensor import TensorType
-from mplang.protos.v1alpha1 import mpir_pb2
+from mplang.v1.core.expr.walk import walk
+from mplang.v1.core.mask import Mask
+from mplang.v1.core.mptype import MPType
+from mplang.v1.core.pfunc import PFunction
+from mplang.v1.core.table import TableType
+from mplang.v1.core.tensor import TensorType
+from mplang.v1.protos.v1alpha1 import mpir_pb2
 
 # Single mapping table for dtype conversion
 DTYPE_MAPPING = {
@@ -204,7 +204,7 @@ def attr_to_proto(py_value: Any) -> mpir_pb2.AttrProto:
             raise TypeError(f"Unsupported tuple/list type: {type(py_value)}")
     elif isinstance(py_value, FuncDefExpr):
         # Convert FuncDefExpr to GraphProto
-        graph = Writer().dumps(py_value)
+        graph = IrWriter().dumps(py_value)
         attr_proto.type = mpir_pb2.AttrProto.GRAPH
         attr_proto.graph.CopyFrom(graph)
     elif isinstance(py_value, PFunction):
@@ -217,7 +217,9 @@ def attr_to_proto(py_value: Any) -> mpir_pb2.AttrProto:
         # Serialize attrs dictionary
         if py_value.attrs:
             for attr_name, attr_value in py_value.attrs.items():
-                attr_proto.func.attrs[attr_name].CopyFrom(attr_to_proto(attr_value))
+                # Skip None-valued attributes to align with top-level attr handling
+                if attr_value is not None:
+                    attr_proto.func.attrs[attr_name].CopyFrom(attr_to_proto(attr_value))
 
         # Note: We don't serialize ins_info and outs_info since they can be
         # inferred from the input expressions during deserialization
@@ -234,7 +236,7 @@ def attr_to_proto(py_value: Any) -> mpir_pb2.AttrProto:
     return attr_proto
 
 
-class Writer:
+class IrWriter:
     """Writer for serializing Expr-based expressions to GraphProto.
 
     This class traverses an expression tree and converts it into a serialized
@@ -491,6 +493,7 @@ class Writer:
             op = self._create_node_proto(expr, "call")
             self._add_single_expr_inputs(op, expr.fn)
             self._add_expr_inputs(op, *expr.args)
+            self._add_attrs(op, name=expr.name)
             self._finalize_node(op, expr)
         elif isinstance(expr, WhileExpr):
             op = self._create_node_proto(expr, "while")
@@ -524,7 +527,7 @@ class Writer:
             raise TypeError(f"Unsupported expr type for serialization: {type(expr)}")
 
 
-class Reader:
+class IrReader:
     """Reader for deserializing GraphProto back to Expr-based expressions.
 
     This class is responsible for converting serialized GraphProto representations
@@ -822,8 +825,12 @@ class Reader:
                 arg_exprs.append(self._value_cache[dep_name])
             else:
                 raise ValueError(f"Input {input_name} not found for call node")
+        # Optional call-site name attribute
+        call_name = None
+        if "name" in node_proto.attrs:
+            call_name = self._proto_to_attr(node_proto.attrs["name"])  # type: ignore[assignment]
 
-        return CallExpr(fn_expr, arg_exprs)
+        return CallExpr(call_name or "", fn_expr, arg_exprs)
 
     def _proto_to_mptype(self, type_proto: mpir_pb2.MPTypeProto) -> MPType:
         """Convert MPTypeProto to MPType."""
@@ -897,7 +904,7 @@ class Reader:
             )
         elif attr_proto.type == mpir_pb2.AttrProto.GRAPH:
             # Handle nested expressions (for control flow)
-            reader = Reader()
+            reader = IrReader()
             return reader.loads(attr_proto.graph)
         else:
             raise TypeError(f"Unsupported attribute type: {attr_proto.type}")

mplang/{core → v1/core}/mpobject.py

@@ -17,14 +17,14 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any
 
-from mplang.core.dtype import DType
-from mplang.core.mask import Mask
-from mplang.core.mptype import MPType
-from mplang.core.table import TableType
-from mplang.core.tensor import Shape
+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.core.cluster import ClusterSpec
+    from mplang.v1.core.cluster import ClusterSpec
 
 
 class MPContext:

mplang/{core → v1/core}/mptype.py

@@ -20,12 +20,12 @@ from typing import TYPE_CHECKING, Any
 import numpy as np
 
 if TYPE_CHECKING:
-    from mplang.core.mpobject import MPObject
+    from mplang.v1.core.mpobject import MPObject
 
-from mplang.core.dtype import STRING, DType
-from mplang.core.mask import Mask
-from mplang.core.table import TableLike, TableType
-from mplang.core.tensor import ScalarType, Shape, TensorLike, TensorType
+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
@@ -195,6 +195,10 @@ class MPType:
         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
@@ -252,9 +256,8 @@ class MPType:
         if not isinstance(other, MPType):
             return False
         return (
-            self._type == other._type
-            and self._pmask == other._pmask
-            and self._attrs == other._attrs
+            self._type == other._type and self._pmask == other._pmask
+            # and self._attrs == other._attrs # TODO(jint): attrs should be optional
         )
 
     def __hash__(self) -> int:
@@ -270,7 +273,7 @@ class MPType:
     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.core.mpobject import MPObject
+        from mplang.v1.core.mpobject import MPObject
 
         if not isinstance(obj, MPObject):
             return False
@@ -373,7 +376,7 @@ class MPType:
                 import pandas as pd
 
                 if isinstance(obj, pd.DataFrame):
-                    from mplang.core.dtype import DType
+                    from mplang.v1.core.dtypes import DType
 
                     schema_dict = {}
                     for col_name in obj.columns:

mplang/{core → v1/core}/pfunc.py

@@ -19,8 +19,8 @@ from collections.abc import Sequence
 from types import MappingProxyType
 from typing import Any
 
-from mplang.core.table import TableType
-from mplang.core.tensor import TensorType
+from mplang.v1.core.table import TableType
+from mplang.v1.core.tensor import TensorType
 
 __all__ = [
     "PFunction",
@@ -33,7 +33,7 @@ class PFunction:
 
     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", "builtin.read")
+    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
@@ -47,7 +47,7 @@ class PFunction:
 
     Args:
         fn_type: The type/category identifier of this PFunction, indicating which
-            backend or handler should process it (e.g., "spu.makeshares", "builtin.read",
+            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)