mplang-nightly 0.1.dev145__py3-none-any.whl → 0.1.dev147__py3-none-any.whl

mplang/core/expr/evaluator.py

@@ -196,6 +196,28 @@ class EvalSemantic:
                     "uniform_cond: predicate is not uniform across parties"
                 )
 
+    # ------------------------------ While helpers ------------------------------
+    def _check_while_predicate(self, cond_result: list[Any]) -> Any:
+        """Validate while_loop predicate evaluation result.
+
+        Ensures the condition function returns exactly one value and that value
+        is non-None. Returns the boolean predicate value for convenience.
+
+        Raises:
+            AssertionError: If condition function returns != 1 value.
+            RuntimeError: If the single predicate value is None.
+        """
+        assert len(cond_result) == 1, (
+            f"Condition function must return a single value, got {cond_result}"
+        )
+        cond_value = cond_result[0]
+        if cond_value is None:
+            raise RuntimeError(
+                "while_loop condition produced None on rank "
+                f"{self.rank}; ensure the predicate yields a boolean for every party."
+            )
+        return cond_value
+
 
 class RecursiveEvaluator(EvalSemantic, ExprVisitor):
     """Recursive visitor-based evaluator."""
@@ -307,12 +329,8 @@ class RecursiveEvaluator(EvalSemantic, ExprVisitor):
             cond_env = dict(zip(expr.cond_fn.params, state, strict=True))
             cond_evaluator = self._fork(cond_env)
             cond_result = expr.cond_fn.body.accept(cond_evaluator)
-
-            assert len(cond_result) == 1, (
-                f"Condition function must return a single value, got {cond_result}"
-            )
-
-            if not cond_result[0]:
+            cond_value = self._check_while_predicate(cond_result)
+            if not cond_value:
                 break
 
             # Call body function with same arguments
@@ -445,8 +463,8 @@ class IterativeEvaluator(EvalSemantic):
                     cond_vals = self._iter_eval_graph(
                         node.cond_fn.body, {**env, **cond_env}
                     )
-                    assert len(cond_vals) == 1
-                    if not bool(cond_vals[0]):
+                    cond_val = self._check_while_predicate(cond_vals)
+                    if not bool(cond_val):
                         break
                     body_env = dict(zip(node.body_fn.params, state, strict=True))
                     new_state = self._iter_eval_graph(

mplang/core/primitive.py

@@ -483,6 +483,20 @@ def uniform_cond(
     if pred_ty.dtype != BOOL:
         raise TypeError(f"uniform_cond predicate must be boolean, got {pred_ty.dtype}")
 
+    # Static pmask rule:
+    # If predicate has a static pmask (not None), it must equal the current trace
+    # context mask. Otherwise some parties would execute a branch without a
+    # defined predicate value (unsafe). To run on a subset either:
+    #   1. Trace the entire uniform_cond under a subset TraceContext (ctx.fork(mask=...))
+    #   2. Broadcast / lift predicate to full mask (e.g. pshfl_s)
+    # Pred pmask None => dynamic: defer to runtime uniformity (if verify_uniform=True).
+    pred_pmask = pred_ty.pmask
+    if pred_pmask is not None and pred_pmask != cur_tracer.mask:
+        raise ValueError(
+            "uniform_cond predicate static pmask mismatch: predicate pmask="
+            f"{pred_pmask} trace mask={cur_tracer.mask}. Trace under a subset "
+            "context (ctx.fork(mask=...)) or broadcast predicate (pshfl_s) to all parties."
+        )
     # Step 1: Trace both branches in separate contexts
     then_tracer = cur_tracer.fork()
     then_tfn = trace(then_tracer, then_fn, *args)
@@ -706,6 +720,22 @@ def while_loop(
             f"Condition function must return a boolean scalar, got dtype {cond_out_var.mptype.dtype}"
         )
 
+    # Static pmask rule:
+    # If the predicate's pmask is statically known it must match the trace context
+    # mask. Otherwise some parties in this context would lack a boolean to drive
+    # control flow (previously could lead to hang via None). To restrict to a subset:
+    #   1. Trace the entire while_loop under a subset context (ctx.fork(mask=submask)), or
+    #   2. Broadcast predicate to full mask (e.g. pshfl_s) before while_loop.
+    # Dynamic predicates (pmask=None) are allowed; runtime guard (evaluator) raises
+    # if any participating party observes None.
+    pred_pmask = cond_out_var.mptype.pmask
+    if pred_pmask is not None and pred_pmask != cur_tracer.mask:
+        raise ValueError(
+            "while_loop predicate static pmask mismatch: predicate pmask="
+            f"{pred_pmask} trace mask={cur_tracer.mask}. Trace under subset context "
+            "or broadcast predicate to all parties."
+        )
+
     # Validate body returns same number of leaves and same dtype/shape per leaf
     if len(body_tfn.out_vars) != len(cond_tfn.in_vars):
         raise ValueError(

{mplang_nightly-0.1.dev145.dist-info → mplang_nightly-0.1.dev147.dist-info}/METADATA

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

{mplang_nightly-0.1.dev145.dist-info → mplang_nightly-0.1.dev147.dist-info}/RECORD

@@ -24,13 +24,13 @@ mplang/core/mpir.py,sha256=V6S9RqegaI0yojhLkHla5nGBi27ASoxlrEs1k4tGubM,37980
 mplang/core/mpobject.py,sha256=0pHSd7SrAFTScCFcB9ziDztElYQn-oIZOKBx47B3QX0,3732
 mplang/core/mptype.py,sha256=09LbMyJp68W0IkbD0s9YLeVssPg3Rl-rcwjTaCfidIQ,15243
 mplang/core/pfunc.py,sha256=PAr8qRhVveWO5HOI0TgdsWjpi4PFi2iEyuTlr9UVKSY,5106
-mplang/core/primitive.py,sha256=PON1EvE8n5hgz1FiKkza7PMWyv9N8Fwjwj8tQag0_LE,38764
+mplang/core/primitive.py,sha256=C1HMbqmkAvLbdgXiHrWPTQ2v2t1uwC_vsGCtI0TEqHY,40574
 mplang/core/table.py,sha256=BqTBZn7Tfwce4vzl3XYhaX5hVmKagVq9-YoERDta6d8,5892
 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=3rtfHVajmCUtkanzgiCWkepRRiTKehPcDF2i5EbcitI,20789
+mplang/core/expr/evaluator.py,sha256=OYmxkr4Lf2qMHnHK-aca-dfMsAAzGRVWuXrxNk_M_3U,21675
 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
@@ -70,8 +70,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.dev145.dist-info/METADATA,sha256=fW2xWPuijyl5cJPFjX1Ibqe1-kyJ86Pg3BYasJ9sqhQ,16547
-mplang_nightly-0.1.dev145.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev145.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev145.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev145.dist-info/RECORD,,
+mplang_nightly-0.1.dev147.dist-info/METADATA,sha256=WxU1KBbckH3JcYQF1jnDD8ph-EE9wKhbompdIg0KTP8,16547
+mplang_nightly-0.1.dev147.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev147.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev147.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev147.dist-info/RECORD,,