mplang-nightly 0.1.dev332__tar.gz → 0.1.dev333__tar.gz

{mplang_nightly-0.1.dev332 → mplang_nightly-0.1.dev333}/PKG-INFO

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

{mplang_nightly-0.1.dev332 → mplang_nightly-0.1.dev333}/mplang/backends/simp_driver/ops.py

@@ -45,8 +45,8 @@ def _wrap_op_as_graph(op: Operation) -> Graph:
     # Create graph inputs
     graph_inputs = [g.add_input(f"in_{i}", any_type) for i in range(len(op.inputs))]
 
-    # Determine output types
-    output_types = [out.type for out in op.outputs] if op.outputs else [any_type]
+    # Determine output types (faithfully mirror the original op)
+    output_types = [out.type for out in op.outputs]
 
     # Add the operation (this handles outputs and value registration)
     g.add_op(
@@ -77,7 +77,7 @@ def _collect_to_hostvars(results: list[Any], num_outputs: int, world_size: int)
         Single DriverVar if num_outputs == 1, else list of DriverVars
     """
     if num_outputs == 0:
-        return None
+        return []
 
     # Transpose [worker][output] -> [output][worker]
     # results[worker_idx] is a list of URIs for that worker's outputs
@@ -117,7 +117,7 @@ def _generic_simp_dispatch(interpreter: Any, op: Operation, *args: Any) -> Any:
     results = driver.collect(futures)
 
     # 4. Assemble into DriverVar(s)
-    num_outputs = len(op.outputs) if op.outputs else 1
+    num_outputs = len(op.outputs)
     return _collect_to_hostvars(results, num_outputs, world_size)
 
 

{mplang_nightly-0.1.dev332 → mplang_nightly-0.1.dev333}/mplang/backends/simp_worker/ops.py

@@ -53,13 +53,19 @@ def _pcall_static_worker_impl(
 
         try:
             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
+            # Use inner graph's actual output count for safety
+            n_out = len(fn_graph.outputs)
+            if n_out == 0:
+                return []
+            return result[0] if n_out == 1 else result
         finally:
             worker.current_parties = prev_parties
     else:
         # No data for this rank
-        return None if len(op.outputs) == 1 else [None] * len(op.outputs)
+        n_out = len(op.outputs)
+        if n_out == 0:
+            return []
+        return None if n_out == 1 else [None] * n_out
 
 
 def _pcall_dynamic_worker_impl(
@@ -68,7 +74,10 @@ def _pcall_dynamic_worker_impl(
     """Worker implementation of pcall_dynamic."""
     fn_graph = op.regions[0]
     result = interpreter.evaluate_graph(fn_graph, list(args))
-    return result[0] if len(op.outputs) == 1 else result
+    n_out = len(fn_graph.outputs)
+    if n_out == 0:
+        return []
+    return result[0] if n_out == 1 else result
 
 
 def _shuffle_static_worker_impl(

{mplang_nightly-0.1.dev332 → mplang_nightly-0.1.dev333}/mplang/runtime/interpreter.py

@@ -1082,7 +1082,9 @@ class Interpreter(AbstractInterpreter):
                     return
 
                 # Store results
-                if len(op.outputs) == 1:
+                if len(op.outputs) == 0:
+                    pass  # Void operation
+                elif len(op.outputs) == 1:
                     env[op.outputs[0]] = result
                 else:
                     for out_val, res in zip(op.outputs, result, strict=True):

{mplang_nightly-0.1.dev332 → mplang_nightly-0.1.dev333}/tests/backends/test_simp_integration.py

@@ -194,6 +194,50 @@ def test_nested_pcall():
         assert values == [20, 40]
 
 
+def test_pcall_static_zero_output():
+    """Regression: pcall_static body returning None must not crash (IndexError).
+
+    When a function returns None, its traced graph has 0 outputs.
+    The driver/worker dispatch must handle this correctly.
+    """
+    sim = simp.make_simulator(world_size=2)
+    with sim:
+
+        def side_effect_only(x):
+            # Produces a traced op but returns nothing
+            run_jax(jnp.abs, x)
+            return None
+
+        x0 = simp.constant((0,), 1)
+        x1 = simp.constant((1,), 2)
+        x_obj = simp.converge(x0, x1)
+
+        # Should not raise IndexError
+        result = pcall_static((0, 1), side_effect_only, x_obj)
+        assert result is None
+
+
+def test_pcall_static_zero_output_compile_evaluate():
+    """Regression: compile+evaluate path with zero-output pcall_static."""
+    sim = simp.make_simulator(world_size=2)
+    with sim:
+
+        def workflow():
+            x0 = simp.constant((0,), 42)
+            x1 = simp.constant((1,), 99)
+            x_obj = simp.converge(x0, x1)
+
+            def body(x):
+                run_jax(jnp.abs, x)
+                return None
+
+            pcall_static((0, 1), body, x_obj)
+
+        traced = mp.compile(workflow)
+        result = mp.evaluate(traced)
+        assert result is None
+
+
 def test_mp_function_decorator():
     """Test @mp.function decorator which implies pcall_static(ALL, ...)."""
     sim = simp.make_simulator(world_size=3)