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

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

@@ -248,6 +248,7 @@ class ClusterSpec:
         world_size: int,
         *,
         endpoints: list[str] | None = None,
+        spu_world_size: int | None = None,
         spu_protocol: str = "SEMI2K",
         spu_field: str = "FM128",
         runtime_version: str = "simulated",
@@ -325,10 +326,14 @@ class ClusterSpec:
 
         # Shared SPU device
         if enable_spu_device:
+            if spu_world_size is None:
+                spu_world_size = world_size
+            spu_members = [nodes[f"node{i}"] for i in range(spu_world_size)]
+
             devices["SP0"] = Device(
                 name="SP0",
                 kind="SPU",
-                members=list(nodes.values()),
+                members=spu_members,
                 config={
                     "protocol": spu_protocol,
                     "field": spu_field,

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

@@ -19,7 +19,7 @@ import threading
 from abc import ABC, abstractmethod
 from typing import Any
 
-from mplang.core.mask import Mask
+from mplang.v1.core.mask import Mask
 
 
 class ICommunicator(ABC):

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

@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     # Imported only for typing to avoid import cycles at runtime.
-    from mplang.core.mpobject import MPContext
+    from mplang.v1.core.mpobject import MPContext
 
 # The global working context.
 _g_ctx: MPContext | None = None

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

@@ -177,6 +177,13 @@ class DType:
             # TypeError if it's not a pandas dtype we can handle
             pass
 
+        try:
+            return cls._from_arrow_dtype(dtype_like)
+        except (ImportError, TypeError):
+            # ImportError if pyarrow is not installed
+            # TypeError if it's not a pyarrow dtype we can handle
+            pass
+
         if isinstance(dtype_like, type) and dtype_like in (bool, int, float, complex):
             return cls.from_python_type(dtype_like)
         elif hasattr(dtype_like, "dtype") and not isinstance(dtype_like, type):
@@ -225,6 +232,37 @@ class DType:
 
         raise TypeError(f"Unsupported pandas dtype: {dtype_like}")
 
+    @classmethod
+    def _from_arrow_dtype(cls, dtype_like: Any) -> DType:
+        try:
+            import pyarrow as pa
+        except ImportError:
+            raise ImportError("pyarrow not available") from None
+
+        if not isinstance(dtype_like, pa.DataType):
+            raise TypeError("Not a pyarrow dtype")
+
+        ARROW_DTYPE_MAPPING = {
+            pa.bool_(): BOOL,
+            pa.int8(): INT8,
+            pa.int16(): INT16,
+            pa.int32(): INT32,
+            pa.int64(): INT64,
+            pa.uint8(): UINT8,
+            pa.uint16(): UINT16,
+            pa.uint32(): UINT32,
+            pa.uint64(): UINT64,
+            pa.float16(): FLOAT16,
+            pa.float32(): FLOAT32,
+            pa.float64(): FLOAT64,
+            pa.string(): STRING,
+            pa.large_string(): STRING,
+        }
+        result = ARROW_DTYPE_MAPPING.get(dtype_like)
+        if result is not None:
+            return result
+        raise TypeError(f"Unsupported arrow dtype: {dtype_like}")
+
     def to_numpy(self) -> np.dtype:
         """Convert custom DType to NumPy dtype."""
         return np.dtype(self.name)

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

@@ -20,7 +20,7 @@ multi-party computation graphs using the visitor pattern.
 """
 
 # Core expression types
-from mplang.core.expr.ast import (
+from mplang.v1.core.expr.ast import (
     AccessExpr,
     CallExpr,
     CondExpr,
@@ -36,12 +36,12 @@ from mplang.core.expr.ast import (
 )
 
 # Built-in evaluator engines
-from mplang.core.expr.evaluator import IEvaluator, create_evaluator
-from mplang.core.expr.printer import Printer
-from mplang.core.expr.transformer import ExprTransformer
+from mplang.v1.core.expr.evaluator import IEvaluator, create_evaluator
+from mplang.v1.core.expr.printer import Printer
+from mplang.v1.core.expr.transformer import ExprTransformer
 
 # Utility functions
-from mplang.core.expr.utils import (
+from mplang.v1.core.expr.utils import (
     deduce_mask,
     ensure_scalar,
     ensure_tensorlist_equal,
@@ -49,8 +49,8 @@ from mplang.core.expr.utils import (
 )
 
 # Visitor pattern interface
-from mplang.core.expr.visitor import ExprVisitor
-from mplang.core.expr.walk import walk, walk_dataflow, walk_structural
+from mplang.v1.core.expr.visitor import ExprVisitor
+from mplang.v1.core.expr.walk import walk, walk_dataflow, walk_structural
 
 __all__ = [
     "AccessExpr",

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

@@ -26,15 +26,15 @@ import logging
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any
 
-from mplang.core.expr.utils import deduce_mask
-from mplang.core.mask import Mask
-from mplang.core.mptype import MPType, Rank
-from mplang.core.pfunc import PFunction
-from mplang.core.table import TableType
-from mplang.core.tensor import TensorType
+from mplang.v1.core.expr.utils import deduce_mask
+from mplang.v1.core.mask import Mask
+from mplang.v1.core.mptype import MPType, Rank
+from mplang.v1.core.pfunc import PFunction
+from mplang.v1.core.table import TableType
+from mplang.v1.core.tensor import TensorType
 
 if TYPE_CHECKING:
-    from mplang.core.expr.visitor import ExprVisitor
+    from mplang.v1.core.expr.visitor import ExprVisitor
 
 
 class Expr(ABC):
@@ -286,8 +286,8 @@ class ConvExpr(Expr):
         # Validate dtype / shape consistency.
         first = types[0]
         for c in types[1:]:
-            if (c.dtype, c.shape) != (first.dtype, first.shape):
-                raise TypeError(f"Inconsistent dtype/shape in pconv: {c} vs {first}")
+            if c.raw_type() != first.raw_type():
+                raise TypeError(f"Inconsistent type in pconv: {c} vs {first}")
 
         # Deduce the pmask by intersecting all pmasks.
         pmasks = [t.pmask for t in types]
@@ -316,7 +316,7 @@ class ConvExpr(Expr):
             else:
                 out_pmask = None
 
-        return [MPType.tensor(first.dtype, first.shape, out_pmask, **first.attrs)]
+        return [MPType(first.raw_type(), out_pmask, first.attrs)]
 
     def accept(self, visitor: ExprVisitor) -> Any:
         return visitor.visit_conv(self)
@@ -398,9 +398,7 @@ class ShflSExpr(Expr):
     def _compute_mptypes(self) -> list[MPType]:
         # The types are the same as the source value, but with a new pmask.
         src_type = self.src_val.mptype
-        return [
-            MPType.tensor(src_type.dtype, src_type.shape, self.pmask, **src_type.attrs)
-        ]
+        return [MPType(src_type._type, self.pmask, src_type.attrs)]
 
     def accept(self, visitor: ExprVisitor) -> Any:
         return visitor.visit_shfl_s(self)

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

@@ -27,8 +27,8 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Any, Protocol
 
-from mplang.core.comm import ICommunicator
-from mplang.core.expr.ast import (
+from mplang.v1.core.comm import ICommunicator
+from mplang.v1.core.expr.ast import (
     AccessExpr,
     CallExpr,
     CondExpr,
@@ -42,12 +42,12 @@ from mplang.core.expr.ast import (
     VariableExpr,
     WhileExpr,
 )
-from mplang.core.expr.visitor import ExprVisitor
-from mplang.core.expr.walk import walk_dataflow
-from mplang.core.mask import Mask
-from mplang.core.pfunc import PFunction
-from mplang.kernels.context import RuntimeContext
-from mplang.kernels.value import Value
+from mplang.v1.core.expr.visitor import ExprVisitor
+from mplang.v1.core.expr.walk import walk_dataflow
+from mplang.v1.core.mask import Mask
+from mplang.v1.core.pfunc import PFunction
+from mplang.v1.kernels.context import RuntimeContext
+from mplang.v1.kernels.value import Value
 
 
 class IEvaluator(Protocol):

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

@@ -20,8 +20,8 @@ from __future__ import annotations
 
 from typing import Any
 
-from mplang.core.dtypes 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):

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):

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.dtypes 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 = {
@@ -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

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.dtypes 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.dtypes 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.dtypes 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",

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

@@ -28,9 +28,9 @@ from typing import Any, ParamSpec, TypeVar, cast
 
 from jax.tree_util import tree_map
 
-from mplang.core.context_mgr import cur_ctx
-from mplang.core.dtypes import BOOL
-from mplang.core.expr.ast import (
+from mplang.v1.core.context_mgr import cur_ctx
+from mplang.v1.core.dtypes import BOOL
+from mplang.v1.core.expr.ast import (
     AccessExpr,
     CallExpr,
     CondExpr,
@@ -40,13 +40,13 @@ from mplang.core.expr.ast import (
     ShflSExpr,
     WhileExpr,
 )
-from mplang.core.interp import InterpContext, InterpVar, apply
-from mplang.core.mask import Mask
-from mplang.core.mpobject import MPContext, MPObject
-from mplang.core.mptype import Rank
-from mplang.core.pfunc import PFunction
-from mplang.core.tracer import TraceContext, TraceVar, trace
-from mplang.utils.func_utils import var_demorph, var_morph
+from mplang.v1.core.interp import InterpContext, InterpVar, apply
+from mplang.v1.core.mask import Mask
+from mplang.v1.core.mpobject import MPContext, MPObject
+from mplang.v1.core.mptype import Rank
+from mplang.v1.core.pfunc import PFunction
+from mplang.v1.core.tracer import TraceContext, TraceVar, trace
+from mplang.v1.utils.func_utils import var_demorph, var_morph
 
 
 def _switch_ctx(ctx: MPContext, obj: MPObject | Any) -> MPObject | Any:
@@ -298,7 +298,7 @@ def uniform_cond(
 
     1. ``pred`` is a boolean scalar whose runtime value is identical for every enabled party.
     2. At least one branch contains multi-party primitives (``seal`` / ``reveal`` /
-       ``srun`` / ``pshfl`` / mask transformations) whose cost or side-effects you
+       ``srun_jax`` / ``pshfl`` / mask transformations) whose cost or side-effects you
        want to avoid if the branch is not taken.
     3. You require the semantic guarantee that the *non-selected* branch does **not**
        perform communication, allocate intermediate buffers, or leak timing/side-effects.
@@ -559,7 +559,7 @@ def while_loop(
         secret-shared reduction).
 
         cond_fn::
-            sealed_sum = smpc.reveal(smpc.srun(lambda x: jnp.sum(x))(smpc.seal(x)))
+            sealed_sum = smpc.reveal(smpc.srun_jax(lambda x: jnp.sum(x), smpc.seal(x)))
             return sealed_sum < constant(10)
 
         body_fn::

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

@@ -18,16 +18,16 @@ from collections.abc import Iterator
 from dataclasses import dataclass, field
 from typing import Any, Protocol, runtime_checkable
 
-from mplang.core.dtypes import DType
+from mplang.v1.core.dtypes import DType
 
 __all__ = ["TableLike", "TableType"]
 
 
 @runtime_checkable
-class TableLike(Protocol):
+class PandasTableLike(Protocol):
     """
     Protocol for objects structurally resembling tables from common libraries
-    (pandas DataFrame, pyarrow Table, etc.), focusing on dtypes and columns attributes.
+    (pandas DataFrame, polars DataFrame, etc.), focusing on dtypes and columns attributes.
     """
 
     @property
@@ -37,6 +37,26 @@ class TableLike(Protocol):
     def columns(self) -> Any: ...
 
 
+@runtime_checkable
+class ArrowSchema(Protocol):
+    @property
+    def names(self) -> list[str]: ...
+    @property
+    def types(self) -> list[Any]: ...
+
+
+@runtime_checkable
+class ArrowTableLike(Protocol):
+    @property
+    def column_names(self) -> list[str]: ...
+
+    @property
+    def schema(self) -> ArrowSchema: ...
+
+
+TableLike = PandasTableLike | ArrowTableLike
+
+
 @dataclass(frozen=True)
 class TableType:
     """Table schema: ordered list of column name-type pairs.
@@ -109,11 +129,19 @@ class TableType:
         Returns:
             TableType instance
         """
-        columns = [
-            (name, DType.from_any(dtype))
-            for name, dtype in zip(table.columns, table.dtypes, strict=True)
-        ]
-        return cls(tuple(columns))
+        if isinstance(table, PandasTableLike):
+            columns = [
+                (name, DType.from_any(dtype))
+                for name, dtype in zip(table.columns, table.dtypes, strict=True)
+            ]
+            return cls(tuple(columns))
+        elif isinstance(table, ArrowTableLike):
+            schema = table.schema
+            columns = [
+                (name, DType.from_any(dtype))
+                for name, dtype in zip(schema.names, schema.types, strict=True)
+            ]
+            return cls(tuple(columns))
 
     def column_names(self) -> tuple[str, ...]:
         """Get all column names."""

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

@@ -19,7 +19,7 @@ from typing import Any, Protocol, runtime_checkable
 
 import numpy as np
 
-from mplang.core.dtypes import DType
+from mplang.v1.core.dtypes import DType
 
 # basic type aliases
 Shape = tuple[int, ...]

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

@@ -60,15 +60,15 @@ from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any, cast
 
-from mplang.core.cluster import ClusterSpec
-from mplang.core.context_mgr import with_ctx
-from mplang.core.expr.ast import Expr, FuncDefExpr, TupleExpr, VariableExpr
-from mplang.core.expr.printer import Printer
-from mplang.core.mask import Mask
-from mplang.core.mpobject import MPContext, MPObject
-from mplang.core.mptype import MPType
-from mplang.core.pfunc import get_fn_name
-from mplang.utils.func_utils import MorphStruct, var_demorph, var_morph
+from mplang.v1.core.cluster import ClusterSpec
+from mplang.v1.core.context_mgr import with_ctx
+from mplang.v1.core.expr.ast import Expr, FuncDefExpr, TupleExpr, VariableExpr
+from mplang.v1.core.expr.printer import Printer
+from mplang.v1.core.mask import Mask
+from mplang.v1.core.mpobject import MPContext, MPObject
+from mplang.v1.core.mptype import MPType
+from mplang.v1.core.pfunc import get_fn_name
+from mplang.v1.utils.func_utils import MorphStruct, var_demorph, var_morph
 
 
 class VarNamer:

mplang/{host.py → v1/host.py}

@@ -19,7 +19,7 @@ from typing import Any
 
 from jax.tree_util import tree_map
 
-from mplang.core import (
+from mplang.v1.core import (
     ClusterSpec,
     InterpContext,
     MPContext,
@@ -28,7 +28,7 @@ from mplang.core import (
     TracedFunction,
     trace,
 )
-from mplang.core.context_mgr import cur_ctx, with_ctx
+from mplang.v1.core.context_mgr import cur_ctx, with_ctx
 
 
 def evaluate(
@@ -76,11 +76,11 @@ def fetch(interp: InterpContext | None, objs: Any) -> Any:  # type: ignore[misc]
     evaluated = evaluate(ctx, lambda x: x, objs)
 
     def fetch_impl(arg: MPObject | Any) -> Any:
-        if isinstance(arg, MPObject):
-            return ctx.fetch(arg)
-        else:
+        if not isinstance(arg, MPObject):
             return arg
 
+        return ctx.fetch(arg)
+
     return tree_map(fetch_impl, evaluated)
 
 

mplang/{kernels → v1/kernels}/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from mplang.kernels.value import (
+from mplang.v1.kernels.value import (
     BytesBlob,
     TableValue,
     TensorValue,

mplang/{kernels → v1/kernels}/base.py

@@ -37,7 +37,7 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-    from mplang.kernels.context import RuntimeContext
+    from mplang.v1.kernels.context import RuntimeContext
 
 __all__ = [
     "KernelContext",