mplang-nightly 0.1.dev167__py3-none-any.whl → 0.1.dev169__py3-none-any.whl

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/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/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/ops/ibis_cc.py

@@ -95,8 +95,8 @@ def is_ibis_function(func: Callable) -> bool:
 _IBIS_MOD = stateless_mod("ibis")
 
 
-class IbisCompiler(FeOperation):
-    """Ibis compiler frontend operation."""
+class IbisRunner(FeOperation):
+    """Ibis runner frontend operation."""
 
     def trace(
         self, func: Callable, *args: Any, **kwargs: Any
@@ -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/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.
@@ -162,21 +162,21 @@ class JaxCompiler(FeOperation):
         along with variable arguments for evaluation.
 
         Args:
-            func: The JAX function to compile
+            jax_fn: The JAX function to compile
             *args: Positional arguments to the function
             **kwargs: Keyword arguments to the function
 
         Returns:
-            tuple[PFunction, list[MPObject], Any]: The compiled PFunction, input variables, and output tree
+            tuple[PFunction, list[MPObject], PyTreeDef]: The compiled PFunction, input variables, and output tree
         """
 
         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/ops/{sql.py → 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/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.dist-info → mplang_nightly-0.1.dev169.dist-info}/METADATA

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

{mplang_nightly-0.1.dev167.dist-info → mplang_nightly-0.1.dev169.dist-info}/RECORD

@@ -1,6 +1,6 @@
 mplang/__init__.py,sha256=ofO-F-CNoVIxpMpTJtTJoQtKegJcHwcOJLzoVispiyc,1852
 mplang/api.py,sha256=ssmv0_CyZPFORhOUJ84Jo6NwRJSK7_Ono3n7ZjEg4sA,3058
-mplang/device.py,sha256=oFdu_XGaJ7AaSePCg1IEQmt6pVoJ3YWHzFyFX4UQBGM,12552
+mplang/device.py,sha256=7X2rPp3hnHFVAYrvBLqdB5XEbURd3VGzXWbw-6bJ1sU,12488
 mplang/analysis/__init__.py,sha256=CTHFvRsi-nFngojqjn08UaR3RY9i7CJ7T2UdR95kCrk,1056
 mplang/analysis/diagram.py,sha256=ffwgD12gL1_KH1uJ_EYkjmIlDrfxYJJkWj-wHl09_Xk,19520
 mplang/core/__init__.py,sha256=lWxlEKfRwX7FNDzgyKZ1fiDMaCiqkyg0j5mKlZD_v7g,2244
@@ -20,7 +20,7 @@ mplang/core/tensor.py,sha256=86u6DogSZMoL0w5XjtTmQm2PhA_VjwybN1b6U4Zzphg,2361
 mplang/core/tracer.py,sha256=dVMfUeCMmPz4o6tLXewGMW1Kpy5gpZORvr9w4MhwDtM,14288
 mplang/core/expr/__init__.py,sha256=qwiSTUOcanFJLyK8HZ13_L1ZDrybqpPXIlTHAyeumE8,1988
 mplang/core/expr/ast.py,sha256=KE46KTtlH9RA2V_EzWVKCKolsycgTmt7SotUrOc8Qxs,20923
-mplang/core/expr/evaluator.py,sha256=91KCnRczzp2vfMevrWs0xpkq0C-jQ_sQG2LOWFvxCUY,23280
+mplang/core/expr/evaluator.py,sha256=EFy71vYUL2xLHCtMkWlYJpyGyujDdVSAx8ByET-62qQ,23297
 mplang/core/expr/printer.py,sha256=VblKGnO0OUfzH7EBkszwRNxQUB8QyyC7BlJWJEUv9so,9546
 mplang/core/expr/transformer.py,sha256=TyL-8FjrVvDq_C9X7kAuKkiqt2XdZM-okjzVQj0A33s,4893
 mplang/core/expr/utils.py,sha256=VDTJ_-CsdHtVy9wDaGa7XdFxQ7o5lYYaeqcgsAhkbNI,2625
@@ -37,15 +37,15 @@ mplang/kernels/spu.py,sha256=vh-WG-uDVPvK11CDzc-f58sUalGts_eWJPjeuxLANfY,12508
 mplang/kernels/sql_duckdb.py,sha256=vq4UCth_PCsH8dxcpx7mMnsq54KtjD8kxwISH7tj3qg,1631
 mplang/kernels/stablehlo.py,sha256=gQyg-2ANyA1TjRg90MZ79mn4cHoXhU7g5GFUCuYXyKs,3231
 mplang/kernels/value.py,sha256=2NZ0UvaLyRYb3bTCqL_fQXMzbHk1qbQ3j9xiv4v5E0A,20726
-mplang/ops/__init__.py,sha256=dpe7WWiYapOFzJeGoKFYBr5mnd6P5SdOyvdYaM2Nhm0,1408
+mplang/ops/__init__.py,sha256=ryPlNziAfc4rwVazdJ6VG7kuUD0bqlRZ2smY5HNIRGI,1018
 mplang/ops/base.py,sha256=h67_SHWNZGUuTCuMll-9kDgGvlPhlFov7WAQCHTmUvw,18258
 mplang/ops/builtin.py,sha256=D7T8rRF9g05VIw9T72lsncF5cDQqaT37eapBieRKvRI,9363
 mplang/ops/crypto.py,sha256=9CeFJrYmvjmgx-3WQl6jHXh8VafRpT4QBunbzsPF8Uc,3646
-mplang/ops/ibis_cc.py,sha256=G8TobBlkyT_IxOWXmMm4APPUnj6sH9bJ-1STfw4kIcw,4256
-mplang/ops/jax_cc.py,sha256=42czYg3hNQbI_nUebXnshlU8ULwM-oBDe_TQoApLNVA,7802
+mplang/ops/ibis_cc.py,sha256=a5OqZVRZ1NzugQPYigdlJcGKbMZHqKh1xkiJen-LtCU,4242
+mplang/ops/jax_cc.py,sha256=kVhJM8i8oPd-yPqyeaZ1hfVxcZPzNhTwjhltDh50hyY,7809
 mplang/ops/phe.py,sha256=SatswExjZWPed8y3qA33BCwIWbvsgHCuCAz_pv2RLLw,6790
 mplang/ops/spu.py,sha256=UHr5DSoqG08xDYER_11OsMVjGGNXXxsvkFoVvXU8uik,4989
-mplang/ops/sql.py,sha256=HyY2i5aGC5W7r62JryFSjQCUDXH3kQz82YADwn4z5uc,2015
+mplang/ops/sql_cc.py,sha256=-9uf75gOxLQlFiKjDm75qIx8Gbun7unvkOxezdSLGlE,2112
 mplang/ops/tee.py,sha256=bOpS_BXG12D6bONikzdF2yt0oVZj9Jyd0g_3IXP8VgE,1281
 mplang/protos/v1alpha1/mpir_pb2.py,sha256=Bros37t-4LMJbuUYVSM65rImUYTtZDhNTIADGbZCKp0,7522
 mplang/protos/v1alpha1/mpir_pb2.pyi,sha256=dLxAtFW7mgFR-HYxC4ExI4jbtEWUGTKBvcKhI3BJ8m0,20972
@@ -64,7 +64,7 @@ mplang/runtime/link_comm.py,sha256=ZHNcis8QDu2rcyyF3rhpxMiJDkczoMA_c0iZ2GDW_bA,2
 mplang/runtime/server.py,sha256=CdmBmpbylEl7XeZj26i0rUmTrPTvl2CVdRgbtR02gcg,16543
 mplang/runtime/session.py,sha256=I2711V-pPRCYibNgBhjboUURdubnL6ltCoh5RvFVabs,10641
 mplang/runtime/simulation.py,sha256=1I_8dIqxivxtYnQK0ofz0oXk3lXXh3-zN0lmNFnWucA,11615
-mplang/simp/__init__.py,sha256=dMd2fALHo7wXhGktvg-S_ovkDd9Bu3gaZnEpsUymKjc,11589
+mplang/simp/__init__.py,sha256=U-JUvaO3EEg0cvEtU1Zihh674vJKQBCDy19gI9t1f-0,11581
 mplang/simp/mpi.py,sha256=Wv_Q16TQ3rdLam6OzqXiefIGSMmagGkso09ycyOkHEs,4774
 mplang/simp/random.py,sha256=7PVgWNL1j7Sf3MqT5PRiWplUu-0dyhF3Ub566iqX86M,3898
 mplang/simp/smpc.py,sha256=tdH54aU4T-GIDPhpmf9NCeJC0G67PdOYc04cyUkOnwE,7119
@@ -73,8 +73,8 @@ mplang/utils/crypto.py,sha256=rvPomBFtznRHc3RPi6Aip9lsU8zW2oxBqGv1K3vn7Rs,1052
 mplang/utils/func_utils.py,sha256=vCJcZmu0bEbqhOQKdpttV2_MBllIcPSN0b8U4WjNGGo,5164
 mplang/utils/spu_utils.py,sha256=S3L9RBkBe2AvSuMSQQ12cBY5Y1NPthubvErSX_7nj1A,4158
 mplang/utils/table_utils.py,sha256=aC-IZOKkSmFkpr3NZchLM0Wt0GOn-rg_xHBHREWBwAU,2202
-mplang_nightly-0.1.dev167.dist-info/METADATA,sha256=BQ-UGDTpsXjgu6txqY7TI7oDsC1bFma8PZ4GxAGKua0,16547
-mplang_nightly-0.1.dev167.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev167.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev167.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev167.dist-info/RECORD,,
+mplang_nightly-0.1.dev169.dist-info/METADATA,sha256=3Ml9Mvi3n9iBnvcVp7dp7lFJABk1hUoIP5AB5BbeQFE,16547
+mplang_nightly-0.1.dev169.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev169.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev169.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev169.dist-info/RECORD,,