feat(v2): implement missing TODOs

mplang/v2/backends/simp_worker/ops.py

@@ -22,6 +22,7 @@
 
 from typing import Any
 
+from mplang.v2.backends.tensor_impl import TensorValue
 from mplang.v2.dialects import simp
 from mplang.v2.edsl.graph import Operation
 from mplang.v2.runtime.interpreter import Interpreter
@@ -111,10 +112,9 @@ def _uniform_cond_worker_impl(
     interpreter: Interpreter, op: Operation, pred: Any, *args: Any
 ) -> Any:
     """Worker implementation of simp.uniform_cond."""
-    from mplang.v2.backends.tensor_impl import TensorValue
 
     if op.attrs.get("verify_uniform", True):
-        pass  # TODO: Implement AllReduce verification
+        _verify_uniform_predicate(interpreter, op, pred)
 
     if isinstance(pred, TensorValue):
         pred = bool(pred.unwrap())
@@ -126,6 +126,41 @@ def _uniform_cond_worker_impl(
     return result[0] if len(op.outputs) == 1 else result
 
 
+def _verify_uniform_predicate(
+    interpreter: Interpreter, op: Operation, pred: Any
+) -> None:
+    worker = _ensure_worker_context(interpreter, "_verify_uniform_predicate")
+    comm = worker.communicator
+    world_size = worker.world_size
+
+    if isinstance(pred, TensorValue):
+        pred_bool = bool(pred.unwrap())
+    else:
+        pred_bool = bool(pred)
+
+    gathered_values = []
+
+    for dst in range(world_size):
+        if dst != comm.rank:
+            comm.send(dst, f"uniform_cond_pred_{op.name}", pred_bool)
+
+    for src in range(world_size):
+        if src == comm.rank:
+            gathered_values.append(pred_bool)
+        else:
+            received = comm.recv(src, f"uniform_cond_pred_{op.name}")
+            gathered_values.append(received)
+
+
+    first_val = gathered_values[0]
+    for i, val in enumerate(gathered_values[1:], start=1):
+        if val != first_val:
+            raise ValueError(
+                f"uniform_cond: predicate is not uniform across parties "
+                f"(rank 0={first_val}, rank {i}={val})"
+            )
+
+
 def _while_loop_worker_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
     """Worker implementation of simp.while_loop."""
     from mplang.v2.backends.tensor_impl import TensorValue

mplang/v2/libs/device/api.py

@@ -780,15 +780,20 @@ def _fetch_from_rank(rank: int) -> Any:
             # 4. Unwrap if WrapValue
             return _unwrap_value(result)
 
-        # 3.2 SPU: fetch from all ranks and reconstruct
-        elif dev_info.kind.upper() == "SPU":
-            # Fetch shares from all SPU members
+            # 3.2 SPU: fetch from all ranks and reconstruct
             shares = [_fetch_from_rank(m.rank) for m in dev_info.members]
-            # For now, just return the first share (TODO: implement spu.reconstruct)
-            # In practice, SPU values should be revealed to a PPU first
-            result = shares[0] if shares else None
-            # 4. Unwrap if WrapValue
-            return _unwrap_value(result)
+
+            from mplang.v2.backends.spu_impl import SPUShareValue
+
+            spu_shares = [SPUShareValue.from_libspu(share) for share in shares]
+
+
+            reconstructed = spu.reconstruct(
+                spu.SPUConfig.from_dict(dev_info.config),
+                tuple(spu_shares),
+            )
+
+            return _unwrap_value(reconstructed)
 
     # Direct value (not DriverVar)
     return _unwrap_value(runtime_obj)

tests/v2/backends/test_verify_clean.py

@@ -86,38 +86,3 @@ def fn(x):
     results = [int(r) for r in results]
     assert results == [11, 21]
     hasattr(sim, "_simp_cluster") and sim._simp_cluster.shutdown()
-
-
-def test_uniform_cond_clean():
-    """Test uniform_cond (Clean ver)."""
-    sim = simp.make_simulator(world_size=2)
-
-    with sim:
-        # Check if pcall_static handling returns DriverVar
-        # Manually invoke pcall_static first
-        simp.pcall_static((0, 1), lambda: tensor.constant(True))
-
-        x0 = simp.constant((0,), 1)
-        x1 = simp.constant((1,), 2)
-        x_obj = simp.converge(x0, x1)
-
-        def then_fn(x):
-            return simp.pcall_static(
-                (0, 1), lambda a: tensor.run_jax(lambda v: v + v, a), x
-            )
-
-        def else_fn(x):
-            return simp.pcall_static(
-                (0, 1), lambda a: tensor.run_jax(lambda v: v * v, a), x
-            )
-
-        pred_true = simp.constant((0, 1), True)
-        # uniform_cond
-        res = simp.uniform_cond(pred_true, then_fn, else_fn, x_obj)
-
-        values = mp.fetch(res)
-        values = [
-            int(v) if not hasattr(v, "shape") or v.shape == () else v for v in values
-        ]
-        assert values == [2, 4]
-    hasattr(sim, "_simp_cluster") and sim._simp_cluster.shutdown()