mplang-nightly 0.1.dev167__tar.gz → 0.1.dev168__tar.gz

{mplang_nightly-0.1.dev167 → mplang_nightly-0.1.dev168}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev167
+Version: 0.1.dev168
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License

{mplang_nightly-0.1.dev167 → mplang_nightly-0.1.dev168}/mplang/core/expr/evaluator.py

@@ -27,8 +27,6 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Any, Protocol
 
-import numpy as np
-
 from mplang.core.comm import ICommunicator
 from mplang.core.expr.ast import (
     AccessExpr,
@@ -234,11 +232,14 @@ class EvalSemantic:
         """Convert Value payloads to numpy/python equivalents when possible."""
         if value is None:
             return None
+
         if isinstance(value, Value):
             # Try to_numpy first for broader compatibility
             to_numpy = getattr(value, "to_numpy", None)
             if callable(to_numpy):
                 arr = to_numpy()
+                import numpy as np
+
                 if isinstance(arr, np.ndarray):
                     if arr.size == 1:
                         return arr.item()

{mplang_nightly-0.1.dev167 → mplang_nightly-0.1.dev168}/mplang/device.py

@@ -37,8 +37,8 @@ from mplang.core.context_mgr import cur_ctx
 from mplang.core.tensor import TensorType
 from mplang.ops import builtin, crypto, ibis_cc, jax_cc, tee
 from mplang.ops.base import FeOperation
-from mplang.ops.ibis_cc import IbisCompiler
-from mplang.ops.jax_cc import JaxCompiler
+from mplang.ops.ibis_cc import IbisRunner
+from mplang.ops.jax_cc import JaxRunner
 from mplang.simp import mpi, smpc
 
 # Automatic transfer between devices when parameter is not on the target device.
@@ -82,7 +82,7 @@ _is_mpobj = lambda x: isinstance(x, MPObject)
 def _device_run_spu(
     dev_info: Device, op: FeOperation, *args: Any, **kwargs: Any
 ) -> Any:
-    if not isinstance(op, JaxCompiler):
+    if not isinstance(op, JaxRunner):
         raise ValueError("SPU device only supports JAX frontend.")
     fn, *aargs = args
     var = smpc.srun(fn)(*aargs, **kwargs)
@@ -92,7 +92,7 @@ def _device_run_spu(
 def _device_run_tee(
     dev_info: Device, op: FeOperation, *args: Any, **kwargs: Any
 ) -> Any:
-    if not isinstance(op, JaxCompiler) and not isinstance(op, IbisCompiler):
+    if not isinstance(op, JaxRunner) and not isinstance(op, IbisRunner):
         raise ValueError("TEE device only supports JAX and Ibis frontend.")
     assert len(dev_info.members) == 1
     rank = dev_info.members[0].rank
@@ -159,11 +159,9 @@ def device(dev_id: str, *, fe_type: str = "jax") -> Callable:
                 return _device_run(dev_id, fn, *args, **kwargs)
             else:
                 if fe_type == "jax":
-                    return _device_run(dev_id, jax_cc.jax_compile, fn, *args, **kwargs)
+                    return _device_run(dev_id, jax_cc.run_jax, fn, *args, **kwargs)
                 elif fe_type == "ibis":
-                    return _device_run(
-                        dev_id, ibis_cc.ibis_compile, fn, *args, **kwargs
-                    )
+                    return _device_run(dev_id, ibis_cc.run_ibis, fn, *args, **kwargs)
                 else:
                     raise ValueError(f"Unsupported frontend type: {fe_type}")
 

{mplang_nightly-0.1.dev167 → mplang_nightly-0.1.dev168}/mplang/ops/__init__.py

@@ -19,28 +19,17 @@ This module contains compilers that transform high-level functions into
 portable, serializable intermediate representations.
 """
 
-from mplang.ops import builtin as builtin
-from mplang.ops import crypto as crypto
-from mplang.ops import ibis_cc as ibis_cc
-from mplang.ops import jax_cc as jax_cc
-from mplang.ops import phe as phe
-from mplang.ops import spu as spu
-from mplang.ops import tee as tee
+from mplang.ops import builtin, crypto, ibis_cc, jax_cc, phe, spu, sql_cc, tee
 from mplang.ops.base import FeOperation as FeOperation
-from mplang.ops.ibis_cc import ibis_compile as ibis_compile
-from mplang.ops.jax_cc import jax_compile as jax_compile
-from mplang.ops.sql import sql_run as sql_run
 
 __all__ = [
     "FeOperation",
     "builtin",
     "crypto",
     "ibis_cc",
-    "ibis_compile",
     "jax_cc",
-    "jax_compile",
     "phe",
     "spu",
-    "sql_run",
+    "sql_cc",
     "tee",
 ]

{mplang_nightly-0.1.dev167 → mplang_nightly-0.1.dev168}/mplang/ops/ibis_cc.py

@@ -95,7 +95,7 @@ def is_ibis_function(func: Callable) -> bool:
 _IBIS_MOD = stateless_mod("ibis")
 
 
-class IbisCompiler(FeOperation):
+class IbisRunner(FeOperation):
     """Ibis compiler frontend operation."""
 
     def trace(
@@ -136,4 +136,4 @@ class IbisCompiler(FeOperation):
         return pfunc, in_vars, treedef
 
 
-ibis_compile = IbisCompiler(_IBIS_MOD, "compile")
+run_ibis = IbisRunner(_IBIS_MOD, "run")

{mplang_nightly-0.1.dev167 → mplang_nightly-0.1.dev168}/mplang/ops/jax_cc.py

@@ -149,11 +149,11 @@ def jax2stablehlo(
     return pfn, in_vars, out_tree
 
 
-class JaxCompiler(FeOperation):
-    """JAX compiler frontend operation."""
+class JaxRunner(FeOperation):
+    """JAX function runner frontend operation."""
 
     def trace(
-        self, func: Callable, *args: Any, **kwargs: Any
+        self, jax_fn: Callable, *args: Any, **kwargs: Any
     ) -> tuple[PFunction, list[MPObject], PyTreeDef]:
         """
         JAX compilation helper function.
@@ -173,10 +173,10 @@ class JaxCompiler(FeOperation):
         def is_variable(arg: Any) -> bool:
             return isinstance(arg, MPObject)
 
-        pfunc, in_vars, out_tree = jax2stablehlo(is_variable, func, *args, **kwargs)
+        pfunc, in_vars, out_tree = jax2stablehlo(is_variable, jax_fn, *args, **kwargs)
         return pfunc, in_vars, out_tree
 
 
 _JAX_MOD = stateless_mod("jax")
 
-jax_compile = JaxCompiler(_JAX_MOD, "compile")
+run_jax = JaxRunner(_JAX_MOD, "run")

mplang_nightly-0.1.dev167/mplang/ops/sql.py → mplang_nightly-0.1.dev168/mplang/ops/sql_cc.py

@@ -22,15 +22,16 @@ from mplang.ops.base import FeOperation, stateless_mod
 _SQL_MOD = stateless_mod("sql")
 
 
-class SqlFE(FeOperation):
+class SqlRunner(FeOperation):
     def __init__(self, dialect: str = "duckdb"):
         # Bind to sql module with a stable op name for registry/dispatch
         super().__init__(_SQL_MOD, "run")
         self._dialect = dialect
 
+    # TODO(jint): we should deduce out_type according to query and in_tables' schema
     def trace(
         self,
-        sql: str,
+        query: str,
         out_type: TableType,
         in_tables: dict[str, MPObject] | None = None,
     ) -> tuple[PFunction, list[MPObject], PyTreeDef]:
@@ -48,7 +49,7 @@ class SqlFE(FeOperation):
         pfn = PFunction(
             fn_type="sql.run",
             fn_name="",
-            fn_text=sql,
+            fn_text=query,
             ins_info=tuple(ins_info),
             outs_info=(out_type,),
             in_names=tuple(in_names),
@@ -58,4 +59,4 @@ class SqlFE(FeOperation):
         return pfn, in_vars, treedef
 
 
-sql_run = SqlFE("duckdb")
+run_sql = SqlRunner("duckdb")

{mplang_nightly-0.1.dev167 → mplang_nightly-0.1.dev168}/mplang/simp/__init__.py

@@ -139,10 +139,10 @@ def run_impl(
         pfunc, eval_args, out_tree = func(*args, **kwargs)
     else:
         if ibis_cc.is_ibis_function(func):
-            pfunc, eval_args, out_tree = ibis_cc.ibis_compile(func, *args, **kwargs)
+            pfunc, eval_args, out_tree = ibis_cc.run_ibis(func, *args, **kwargs)
         else:
             # unknown python callable, treat it as jax function
-            pfunc, eval_args, out_tree = jax_cc.jax_compile(func, *args, **kwargs)
+            pfunc, eval_args, out_tree = jax_cc.run_jax(func, *args, **kwargs)
     results = peval(pfunc, eval_args, pmask)
     return out_tree.unflatten(results)
 

{mplang_nightly-0.1.dev167 → mplang_nightly-0.1.dev168}/tests/ops/test_sql.py

@@ -18,14 +18,14 @@ import mplang
 from mplang import simp
 from mplang.core.dtype import INT32
 from mplang.core.table import TableType
-from mplang.ops import sql_run
+from mplang.ops import sql_cc
 
 
 def test_sqlrun():
     # test without input
     sql = "select a from 'a.csv'"
     out_type = TableType.from_pairs([("a", INT32)])
-    pfn, input_args, _output_tree = sql_run(sql, out_type)
+    pfn, input_args, _output_tree = sql_cc.run_sql(sql, out_type)
     assert pfn.fn_text == sql
     assert pfn.outs_info[0] == out_type
     assert len(input_args) == 0
@@ -40,7 +40,7 @@ def test_sqlrun():
     sql = "select a from table"
     out_type = TableType.from_pairs([("a", INT32)])
     in_tables = {"table": in_tbl}
-    pfn, input_args, _output_tree = sql_run(sql, out_type, in_tables=in_tables)
+    pfn, input_args, _output_tree = sql_cc.run_sql(sql, out_type, in_tables=in_tables)
     assert pfn.fn_text == sql
     assert pfn.outs_info[0] == out_type
     assert len(input_args) == 1