mplang-nightly 0.1.dev250__tar.gz → 0.1.dev252__tar.gz

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/PKG-INFO

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

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/examples/v2/pcall_jax.py

@@ -22,9 +22,7 @@ import jax.numpy as jnp
 import numpy as np
 
 import mplang.v2.edsl as el
-import mplang.v2.edsl.typing as elt
-from mplang.v2.dialects.simp import pcall_static
-from mplang.v2.dialects.tensor import jax_fn
+from mplang.v2.dialects import simp, tensor
 
 
 def main():
@@ -37,27 +35,21 @@ def main():
     def add(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
         return jnp.add(x, y)
 
-    # Prepare input data (MP types)
-    x = el.InterpObject(
-        np.array([1.0, 2.0, 3.0]),
-        elt.MPType(elt.TensorType(elt.f32, (3,)), (0,)),  # P0 holds
-    )
-
-    y = el.InterpObject(
-        np.array([4.0, 5.0, 6.0]),
-        elt.MPType(elt.TensorType(elt.f32, (3,)), (0,)),  # P0 holds
-    )
-
     # Define computation using pcall + jax_fn
-    def compute(x, y):
+    # Use simp.constant to create data at parties
+    def compute():
+        # P0 holds x and y as constants
+        x = simp.constant((0,), np.array([1.0, 2.0, 3.0], dtype=np.float32))
+        y = simp.constant((0,), np.array([4.0, 5.0, 6.0], dtype=np.float32))
+
         # P0 executes JAX square computation
-        squared = pcall_static((0,), jax_fn(square), x)
+        squared = simp.pcall_static((0,), tensor.jax_fn(square), x)
         # P0 executes JAX add computation
-        result = pcall_static((0,), jax_fn(add), squared, y)
+        result = simp.pcall_static((0,), tensor.jax_fn(add), squared, y)
         return result
 
     # Trace to generate graph
-    traced = el.trace(compute, x, y)
+    traced = el.trace(compute)
 
     # Print computation graph using EDSL printer
     print("=" * 70)

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/examples/v2/stax_nn.py

@@ -163,9 +163,9 @@ def train(
 
     opt_state_p0_final = do_train(opt_state)
 
-    # Fetch results to driver
+    # Fetch results to driver (data is already on P0 with device attribute)
     print("Fetching results...")
-    opt_state_final = tree_map(lambda x: mp.fetch(x, party="P0"), opt_state_p0_final)
+    opt_state_final = tree_map(lambda x: mp.fetch(x), opt_state_p0_final)
 
     # Extract params
     params = get_params(opt_state_final)

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/__init__.py

@@ -89,6 +89,7 @@ from mplang.v2.libs.device import (
     put,
     set_dev_attr,
 )
+from mplang.v2.libs.device import fetch as device_fetch
 from mplang.v2.runtime.interpreter import Interpreter
 
 # =============================================================================
@@ -186,88 +187,80 @@ def evaluate(
             inputs = fn.prepare_inputs(*args, **kwargs)
             inputs = [unwrap_if_interp(v) for v in inputs]
             raw_result = interp.evaluate_graph(fn.graph, inputs)
-            return fn.reconstruct_outputs(raw_result)
+            wrapped = [
+                InterpObject(v, fn.graph.outputs[i].type, interp)
+                for i, v in enumerate(raw_result)
+            ]
+            return fn.reconstruct_outputs(wrapped)
 
         return fn(*args, **kwargs)
 
 
 def fetch(
     result: Any,
-    party: int | str | None = None,
     *,
+    follow_device: bool = True,
     context: Interpreter | None = None,
 ) -> Any:
-    """Fetch the result, optionally for a specific party.
+    """Fetch results from interpreter context to Python.
 
     Args:
-        result: The result to fetch (DriverVar or other value).
-        party: Optional party index or device name (e.g., "P0").
-        context: Optional interpreter context. If None, uses current context.
+        result: Object(s) to fetch. Can be a single InterpObject, DriverVar,
+            or nested structure containing them.
+        follow_device: If True and object has device attribute, dispatch to
+            device.fetch which fetches from the correct rank based on device.
+            If False, fetch from all parties.
+        context: Interpreter context. If None, uses current context.
 
     Returns:
-        The fetched data (unwrapped from Value wrappers).
-
-    Example:
-        >>> with mp.make_simulator(3) as sim:
-        ...     data = mp.fetch(result, "P0")  # uses sim from context
+        Fetched Python values. For device objects with follow_device=True,
+        returns single value from the device's rank(s). Otherwise returns
+        list of values (one per party) or single value for world_size=1.
     """
-    from typing import cast
 
-    from mplang.v2.backends.simp_driver.state import SimpDriver
+    from jax.tree_util import tree_map
+
     from mplang.v2.backends.simp_driver.values import DriverVar
-    from mplang.v2.backends.table_impl import TableValue
-    from mplang.v2.backends.tensor_impl import TensorValue
     from mplang.v2.runtime.interpreter import InterpObject
+    from mplang.v2.runtime.value import WrapValue
 
     interp = _get_context(context)
 
-    def _unwrap_value(val: Any) -> Any:
-        """Unwrap Value types to get the underlying data."""
-        if isinstance(val, TensorValue):
-            return val.data
-        elif isinstance(val, TableValue):
-            return val.data
-        return val
-
-    # Unwrap InterpObject to get the runtime value
-    if isinstance(result, InterpObject):
-        result = result.runtime_obj
-
-    # Get simp state for fetching
-    simp_state = cast(SimpDriver | None, interp.get_dialect_state("simp"))
-    cluster_spec = getattr(interp, "_cluster_spec", None)
-
-    # Fetch from DriverVar
-    if isinstance(result, DriverVar):
-        resolved_values = []
-        for rank, val in enumerate(result.values):
-            if isinstance(val, str) and "://" in val:
-                if simp_state is not None:
-                    fut = simp_state.fetch(rank, val)
-                    resolved_values.append(fut.result())
-                else:
-                    resolved_values.append(val)
-            else:
-                resolved_values.append(val)
-
-        # Select party if needed
-        if party is not None:
-            if isinstance(party, str) and cluster_spec is not None:
-                device_info = cluster_spec.devices.get(party)
-                if device_info and device_info.members:
-                    party = device_info.members[0].rank
+    def _fetch_single(var: Any) -> Any:
+        """Fetch a single value from InterpObject."""
+        # InterpObject (from mp.evaluate) - extract runtime_obj
+        if isinstance(var, InterpObject):
+            if follow_device and is_device_obj(var):
+                return device_fetch(var)
+            var = var.runtime_obj  # extract and continue processing
+
+        # DriverVar (simp dialect) - remote fetch from workers
+        if isinstance(var, DriverVar):
+            from mplang.v2.backends.simp_driver.state import SimpDriver
+
+            simp_state = interp.get_dialect_state("simp")
+            assert isinstance(simp_state, SimpDriver), "DriverVar requires simp state"
+
+            resolved: list[Any] = []
+            for rank, uri in enumerate(var.values):
+                if uri is None:
+                    resolved.append(None)
                 else:
-                    raise ValueError(f"Unknown party: {party}")
-            else:
-                # Default logic for int
-                pass
+                    fetched = simp_state.fetch(rank, uri).result()
+                    if isinstance(fetched, WrapValue):
+                        fetched = fetched.data
+                    resolved.append(fetched)
+
+            return resolved[0] if len(resolved) == 1 else resolved
+
+        # WrapValue (TensorValue, TableValue, etc.) - unwrap
+        if isinstance(var, WrapValue):
+            return var.data
 
-            p_idx = cast(int, party)
-            return _unwrap_value(resolved_values[p_idx])  # type: ignore[no-any-return]
-        return [_unwrap_value(v) for v in resolved_values]
+        # Plain values pass through
+        return var
 
-    # Unwrap Value types to get the underlying data
-    return _unwrap_value(result)
+    return tree_map(_fetch_single, result)
 
 
 # Alias for compatibility

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/backends/func_impl.py

@@ -102,4 +102,6 @@ def _call_impl(
         raise TypeError(f"func.call expects FunctionValue, got {type(fn_obj)}")
 
     call_args = list(args)
-    return interpreter.evaluate_graph(fn_obj.graph, call_args)
+    result = interpreter.evaluate_graph(fn_obj.graph, call_args)
+    # Return single value or list based on graph outputs
+    return result[0] if len(fn_obj.graph.outputs) == 1 else result

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/backends/simp_driver/mem.py

@@ -201,27 +201,20 @@ class SimpMemDriver(SimpDriver):
         worker_interp = self._workers[rank]
         worker_ctx = cast(SimpWorker, worker_interp.get_dialect_state("simp"))
 
-        # Resolve URI inputs
-        resolved_inputs = []
-        for inp in inputs:
-            if isinstance(inp, str) and "://" in inp:
-                resolved_inputs.append(worker_ctx.store.get(inp))
-            else:
-                resolved_inputs.append(inp)
+        # Resolve URI inputs (None means rank has no data)
+        resolved_inputs = [
+            worker_ctx.store.get(inp) if inp is not None else None for inp in inputs
+        ]
 
         # Execute
         results = worker_interp.evaluate_graph(graph, resolved_inputs, job_id)
 
-        # Store results
+        # Store results (results is always a list)
         if not graph.outputs:
             return None
-        if len(graph.outputs) == 1:
-            val = results
-            if val is None:
-                return None
-            return worker_ctx.store.put(val)
-        else:
-            return [worker_ctx.store.put(res) for res in results]
+        return [
+            worker_ctx.store.put(res) if res is not None else None for res in results
+        ]
 
 
 def make_simulator(

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/backends/simp_driver/ops.py

@@ -68,7 +68,8 @@ def _collect_to_hostvars(results: list[Any], num_outputs: int, world_size: int)
     """Collect worker results into DriverVar(s).
 
     Args:
-        results: List of results from each worker (length = world_size)
+        results: List of results from each worker (length = world_size).
+                 Each result is a list of URIs (one per output).
         num_outputs: Number of outputs per worker
         world_size: Total number of workers
 
@@ -78,15 +79,16 @@ def _collect_to_hostvars(results: list[Any], num_outputs: int, world_size: int)
     if num_outputs == 0:
         return None
 
-    if num_outputs == 1:
-        return DriverVar(results)
-
-    # Multiple outputs: transpose [worker][output] -> [output][worker]
+    # Transpose [worker][output] -> [output][worker]
+    # results[worker_idx] is a list of URIs for that worker's outputs
     transposed = []
-    for i in range(num_outputs):
+    for out_idx in range(num_outputs):
         transposed.append(
-            DriverVar([res[i] if res is not None else None for res in results])
+            DriverVar([res[out_idx] if res is not None else None for res in results])
         )
+
+    if num_outputs == 1:
+        return transposed[0]
     return transposed
 
 

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/backends/simp_worker/http.py

@@ -244,11 +244,18 @@ def create_worker_app(
 
     def _do_execute(graph: Graph, inputs: list[Any], job_id: str | None = None) -> Any:
         """Execute graph in worker thread."""
-        # Note: we ignore job_id for now as generic Interpreter doesn't take it.
-        # But we could set it in thread local if needed.
-        result = worker.evaluate_graph(graph, inputs)
+        # Resolve URI inputs (None means rank has no data)
+        resolved_inputs = [
+            store.get(inp) if inp is not None else None for inp in inputs
+        ]
+
+        result = worker.evaluate_graph(graph, resolved_inputs)
         comm.wait_pending_sends()
-        return result
+
+        # Store results and return URIs (result is always a list)
+        if not graph.outputs:
+            return None
+        return [store.put(res) if res is not None else None for res in result]
 
     @app.post("/exec")
     async def execute(req: ExecRequest) -> dict[str, str]:

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/backends/simp_worker/ops.py

@@ -51,17 +51,17 @@ def _pcall_static_worker_impl(
         interpreter.current_parties = parties  # type: ignore[attr-defined]
 
         try:
-            return interpreter.evaluate_graph(fn_graph, list(args))
+            result = interpreter.evaluate_graph(fn_graph, list(args))
+            # Return single value for single output (interpreter expects this)
+            return result[0] if len(op.outputs) == 1 else result
         finally:
             if prev_parties is None:
                 del interpreter.current_parties  # type: ignore[attr-defined]
             else:
                 interpreter.current_parties = prev_parties  # type: ignore[attr-defined]
     else:
-        if len(op.outputs) == 1:
-            return None
-        else:
-            return [None] * len(op.outputs)
+        # No data for this rank
+        return None if len(op.outputs) == 1 else [None] * len(op.outputs)
 
 
 def _pcall_dynamic_worker_impl(
@@ -69,7 +69,8 @@ def _pcall_dynamic_worker_impl(
 ) -> Any:
     """Worker implementation of pcall_dynamic."""
     fn_graph = op.regions[0]
-    return interpreter.evaluate_graph(fn_graph, list(args))
+    result = interpreter.evaluate_graph(fn_graph, list(args))
+    return result[0] if len(op.outputs) == 1 else result
 
 
 def _shuffle_static_worker_impl(
@@ -122,9 +123,10 @@ def _uniform_cond_worker_impl(
         pred = bool(pred.unwrap())
 
     if pred:
-        return interpreter.evaluate_graph(op.regions[0], list(args))
+        result = interpreter.evaluate_graph(op.regions[0], list(args))
     else:
-        return interpreter.evaluate_graph(op.regions[1], list(args))
+        result = interpreter.evaluate_graph(op.regions[1], list(args))
+    return result[0] if len(op.outputs) == 1 else result
 
 
 def _while_loop_worker_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
@@ -142,22 +144,20 @@ def _while_loop_worker_impl(interpreter: Interpreter, op: Operation, *args: Any)
         region_inputs = current_state + captures
 
         cond_res = interpreter.evaluate_graph(cond_graph, region_inputs)
+        # cond_res is a list, extract the single boolean
+        cond_val = cond_res[0] if cond_res else False
 
-        if isinstance(cond_res, TensorValue):
-            cond_res = bool(cond_res.unwrap())
+        if isinstance(cond_val, TensorValue):
+            cond_val = bool(cond_val.unwrap())
 
-        if not cond_res:
+        if not cond_val:
             break
 
         body_res = interpreter.evaluate_graph(body_graph, region_inputs)
-        if isinstance(body_res, list):
-            current_state = body_res
-        else:
-            current_state = [body_res]
-
-    if len(current_state) == 1:
-        return current_state[0]
-    return current_state
+        current_state = body_res  # body_res is always a list now
+
+    # Return single value for single output
+    return current_state[0] if len(current_state) == 1 else current_state
 
 
 WORKER_HANDLERS = {

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/backends/tensor_impl.py

@@ -266,8 +266,9 @@ def elementwise_impl(interpreter: Interpreter, op: Operation, *args: Value) -> A
     subgraph = op.regions[0]
 
     if shape == ():
-        # Scalar case
-        return interpreter.evaluate_graph(subgraph, list(args))
+        # Scalar case - return first element from result list
+        result = interpreter.evaluate_graph(subgraph, list(args))
+        return result[0] if len(result) == 1 else result
 
     for index in np.ndindex(shape):
         # Prepare inputs for this element (list ordered by subgraph.inputs)
@@ -295,7 +296,10 @@ def elementwise_impl(interpreter: Interpreter, op: Operation, *args: Value) -> A
                     scalar_inputs.append(arg)
 
         # Recursive execution
-        scalar_out = interpreter.evaluate_graph(subgraph, scalar_inputs)
+        scalar_out_list = interpreter.evaluate_graph(subgraph, scalar_inputs)
+        scalar_out = (
+            scalar_out_list[0] if len(scalar_out_list) == 1 else scalar_out_list
+        )
 
         # Unwrap result if it's a TensorValue (to store in numpy array)
         # We store raw values in the object array for now, but will wrap the final array

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/edsl/tracer.py

@@ -573,24 +573,19 @@ class TracedFunction:
         all_inputs = explicit_inputs + list(self.captured)
         return all_inputs
 
-    def reconstruct_outputs(self, execution_result: Any) -> Any:
+    def reconstruct_outputs(self, execution_result: list[Any]) -> Any:
         """Reconstruct structured output from execution result.
 
         Used by the runtime to format the result of graph execution.
 
         Args:
-            execution_result: Raw result from interpreter.evaluate_graph().
+            execution_result: List of results from interpreter.evaluate_graph().
 
         Returns:
             Structured output matching the original function's return signature.
         """
-        # Normalize to list
-        if len(self.graph.outputs) == 1:
-            results = [execution_result]
-        elif len(self.graph.outputs) == 0:
-            results = []
-        else:
-            results = execution_result  # It's already a list/tuple
+        # execution_result is always a list (now that evaluate_graph returns list)
+        results = execution_result
 
         # Reconstruct
         total_len = len(self.out_imms) + len(self.out_var_pos)

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/libs/device/__init__.py

@@ -25,6 +25,7 @@ from .api import (
     DeviceInferenceError,
     DeviceNotFoundError,
     device,
+    fetch,
     get_dev_attr,
     is_device_obj,
     put,
@@ -41,6 +42,7 @@ __all__ = [
     "DeviceNotFoundError",
     "Node",
     "device",
+    "fetch",
     "get_dev_attr",
     "is_device_obj",
     "jax_fn",

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/libs/device/api.py

@@ -715,3 +715,74 @@ def put(to_dev_id: str, obj: Any) -> Object:
 
     else:
         raise DeviceError(f"Cannot put to device kind '{dev_info.kind}'")
+
+
+def fetch(obj: Object) -> Any:
+    """Fetch data from device to host based on device attribute.
+
+    This function fetches data from the device the object resides on.
+    For PPU/TEE: fetches from the single member rank.
+    For SPU: fetches from all parties (returns reconstructed value).
+
+    Args:
+        obj: Object with device attribute to fetch.
+
+    Returns:
+        Python value (numpy array, scalar, etc.)
+    """
+    from mplang.v2.backends.simp_driver.state import SimpDriver
+    from mplang.v2.backends.simp_driver.values import DriverVar
+    from mplang.v2.backends.table_impl import TableValue
+    from mplang.v2.backends.tensor_impl import TensorValue
+    from mplang.v2.edsl.context import get_current_context
+    from mplang.v2.runtime.interpreter import InterpObject, Interpreter
+
+    def _unwrap_value(val: Any) -> Any:
+        """Unwrap Value types to get the underlying data."""
+        if isinstance(val, TensorValue):
+            return val.data
+        elif isinstance(val, TableValue):
+            return val.data
+        return val
+
+    # Get device info
+    if not is_device_obj(obj):
+        raise DeviceError(
+            "Object does not have device attribute. Use mp.fetch() directly."
+        )
+
+    dev_id = get_dev_attr(obj)
+    cluster = _resolve_cluster()
+    dev_info = cluster.devices[dev_id]
+
+    # Get interpreter context
+    ctx = get_current_context()
+    if not isinstance(ctx, Interpreter):
+        raise RuntimeError("No interpreter context available for fetch")
+
+    simp_state = cast(SimpDriver | None, ctx.get_dialect_state("simp"))
+
+    # Unwrap InterpObject
+    assert isinstance(obj, InterpObject), f"Expected InterpObject, got {type(obj)}"
+    runtime_obj = obj.runtime_obj
+
+    def _fetch_from_rank(rank: int) -> Any:
+        """Fetch value from a rank (DriverVar values are always URIs)."""
+        uri = runtime_obj.values[rank]
+        assert isinstance(uri, str) and "://" in uri, f"Expected URI, got {uri}"
+        assert simp_state is not None, "No simp state for fetch"
+        return _unwrap_value(simp_state.fetch(rank, uri).result())
+
+    # Handle DriverVar
+    if isinstance(runtime_obj, DriverVar):
+        # For PPU/TEE: single member
+        if dev_info.kind.upper() in ("PPU", "TEE"):
+            assert len(dev_info.members) == 1
+            return _fetch_from_rank(dev_info.members[0].rank)
+
+        # For SPU: fetch from first member (should be revealed first)
+        elif dev_info.kind.upper() == "SPU":
+            return _fetch_from_rank(dev_info.members[0].rank)
+
+    # Direct value
+    return _unwrap_value(runtime_obj)

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/mplang/v2/runtime/interpreter.py

@@ -442,13 +442,7 @@ class Interpreter(AbstractInterpreter):
 
         # Execute graph (may have 0 outputs if all were immediates)
         if graph.outputs:
-            result_runtime = self.evaluate_graph(graph, inputs_list)
-
-            # Normalize runtime result to list for structural matching
-            if len(graph.outputs) == 1:
-                result_runtime_list = [result_runtime]
-            else:
-                result_runtime_list = list(result_runtime)
+            result_runtime_list = self.evaluate_graph(graph, inputs_list)
         else:
             result_runtime_list = []
 
@@ -557,10 +551,6 @@ class Interpreter(AbstractInterpreter):
                 # Execute graph
                 results_runtime = self.evaluate_graph(graph, inputs_list)
 
-                # If single output, wrap in list for uniform handling
-                if len(target_outputs) == 1:
-                    results_runtime = [results_runtime]
-
                 # Cache all results
                 for val, res in zip(target_outputs, results_runtime, strict=True):
                     # Wrap as InterpObject and cache
@@ -586,10 +576,9 @@ class Interpreter(AbstractInterpreter):
             # Constants: pass through unchanged
             return obj
 
-    # TODO: change output sigature to list[Any]
     def evaluate_graph(
         self, graph: Graph, inputs: list[Any], job_id: str | None = None
-    ) -> Any:
+    ) -> list[Any]:
         """Execute a Graph IR with runtime data.
 
         Can be overridden by subclasses to implement remote execution or compilation.
@@ -600,7 +589,7 @@ class Interpreter(AbstractInterpreter):
             job_id: Optional unique ID for this execution job (for profiling/tracing).
 
         Returns:
-            Runtime execution results corresponding to graph.outputs
+            List of runtime execution results corresponding to graph.outputs.
         """
         if self.executor:
             return self._evaluate_graph_async(graph, inputs, job_id)
@@ -609,7 +598,7 @@ class Interpreter(AbstractInterpreter):
 
     def _evaluate_graph_sync(
         self, graph: Graph, inputs: list[Any], job_id: str | None = None
-    ) -> Any:
+    ) -> list[Any]:
         """Synchronous execution (Baseline)."""
         # Local environment: Value -> Runtime Object
         env = dict(zip(graph.inputs, inputs, strict=True))
@@ -648,7 +637,6 @@ class Interpreter(AbstractInterpreter):
                     f"No implementation registered for opcode: {op.opcode}"
                 )
 
-            # Update environment with outputs
             # Update environment with outputs
             # Handler should return a single value or a tuple/list of values
             if len(op.outputs) == 0:
@@ -667,14 +655,11 @@ class Interpreter(AbstractInterpreter):
         if self.tracer and job_id:
             self.tracer.save_trace(job_id=job_id, rank=self.trace_pid)
 
-        if len(graph.outputs) == 1:
-            return env[graph.outputs[0]]
-        else:
-            return [env[out] for out in graph.outputs]
+        return [env[out] for out in graph.outputs]
 
     def _evaluate_graph_async(
         self, graph: Graph, inputs: list[Any], job_id: str | None = None
-    ) -> Any:
+    ) -> list[Any]:
         """Asynchronous execution with non-blocking DAG scheduling."""
         # Tracer setup (if not provided, use a disabled stub)
         tracer: ExecutionTracer | _NullTracer
@@ -851,5 +836,4 @@ class Interpreter(AbstractInterpreter):
         if self.tracer and job_id:
             self.tracer.save_trace(job_id=job_id, rank=self.trace_pid)
 
-        final_results = [env[out] for out in graph.outputs]
-        return final_results[0] if len(final_results) == 1 else final_results
+        return [env[out] for out in graph.outputs]

{mplang_nightly-0.1.dev250 → mplang_nightly-0.1.dev252}/tests/v2/backends/simp_driver/test_http.py

@@ -161,8 +161,8 @@ class TestDriverExecution:
                 return y
 
             result = mp.evaluate(transfer)
-            # y is on P1, so fetch from P1
-            value = mp.fetch(result, party="P1")
+            # y is on P1, so fetch will use device attribute to get from P1
+            value = mp.fetch(result)
             assert value == 100
 
     @pytest.mark.skip(
@@ -193,7 +193,8 @@ class TestDriverFetch:
                 return mp.device("P0")(lambda: 999)()
 
             result = mp.evaluate(create_on_p0)
-            value = mp.fetch(result, party="P0")
+            # fetch uses device attribute to get from P0
+            value = mp.fetch(result)
             assert value == 999
 
     def test_fetch_all_parties(self, driver_cluster):