mplang-nightly 0.1.dev148__py3-none-any.whl → 0.1.dev149__py3-none-any.whl

mplang/backend/stablehlo.py

@@ -51,8 +51,15 @@ def _stablehlo_exec(pfunc: PFunction, *args: Any) -> Any:
             raise RuntimeError(f"StableHLO compile failed: {e}") from e
         cache[mlir_text] = compiled
 
+    # Handle JAX's unused parameter elimination via arg_keep_map
+    runtime_args = args
+    if "arg_keep_map" in pfunc.attrs:
+        keep_indices = pfunc.attrs["arg_keep_map"]
+        # Filter out arguments that were eliminated by JAX during compilation
+        runtime_args = tuple(args[i] for i in keep_indices)
+
     jax_args = []
-    for arg in args:
+    for arg in runtime_args:
         if hasattr(arg, "numpy"):
             jax_arg = jnp.array(arg.numpy())  # type: ignore
         else:

mplang/frontend/jax_cc.py

@@ -106,14 +106,46 @@ def jax2stablehlo(
     out_info_flat, out_tree = tree_flatten(lowered.out_info)
     out_info_flat = [TensorType.from_obj(info) for info in out_info_flat]
 
+    # Extract argument keep mapping to handle JAX's unused parameter elimination
+    # JAX can eliminate unused parameters during compilation, but the runtime still
+    # receives all original arguments. We need the mapping to filter them correctly.
+    arg_keep_map = None
+    original_arg_count = len(in_vars)
+
+    try:
+        # Access JAX internal kept_var_idx - the authoritative source
+        # This tells us exactly which original parameters survived compilation
+        compile_args = lowered._lowering.compile_args
+        kept_var_idx = compile_args["kept_var_idx"]
+
+        kept_indices = sorted(kept_var_idx)
+        if len(kept_indices) < original_arg_count:
+            arg_keep_map = kept_indices
+
+    except (AttributeError, KeyError, TypeError) as e:
+        # JAX internal API is not available or changed
+        # This is a hard error - we cannot reliably handle unused parameters
+        # without knowing exactly which ones were kept
+        raise RuntimeError(
+            f"Cannot access JAX's kept_var_idx to handle unused parameter elimination. "
+            f"This function may have unused parameters that JAX optimized away, "
+            f"but we cannot determine which ones without the internal API. "
+            f"Original error: {e}"
+        ) from e
+
     # This format tells JaxRT how to handle the compiled result
-    pfn = PFunction(
-        fn_type="mlir.stablehlo",  # Key: specify StableHLO MLIR format
-        ins_info=tuple(TensorType.from_obj(x) for x in in_vars),
-        outs_info=tuple(out_info_flat),
-        fn_name=get_fn_name(flat_fn),
-        fn_text=mlir_text,  # MLIR text, serializable for transmission
-    )
+    pfn_kwargs: dict[str, Any] = {
+        "fn_type": "mlir.stablehlo",  # Key: specify StableHLO MLIR format
+        "ins_info": tuple(TensorType.from_obj(x) for x in in_vars),
+        "outs_info": tuple(out_info_flat),
+        "fn_name": get_fn_name(flat_fn),
+        "fn_text": mlir_text,  # MLIR text, serializable for transmission
+    }
+
+    if arg_keep_map is not None:
+        pfn_kwargs["arg_keep_map"] = arg_keep_map
+
+    pfn = PFunction(**pfn_kwargs)
     return pfn, in_vars, out_tree
 
 

{mplang_nightly-0.1.dev148.dist-info → mplang_nightly-0.1.dev149.dist-info}/METADATA

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

{mplang_nightly-0.1.dev148.dist-info → mplang_nightly-0.1.dev149.dist-info}/RECORD

@@ -11,7 +11,7 @@ mplang/backend/crypto.py,sha256=H_s5HI7lUP7g0xz-a9qMbSn6dhJStUilKbn3-7SIh0I,3812
 mplang/backend/phe.py,sha256=uNqmrbDAbd97TWS_O6D5sopastHy6J20R7knFE4M4uc,65247
 mplang/backend/spu.py,sha256=QT1q5uv-5P_nBGtTvtA_yI2h3h3zIqNSnvzGT7Shua4,9307
 mplang/backend/sql_duckdb.py,sha256=U_KzEUinxrBRDoUz2Vh597-N4I3hPOBT0RT3tX-ZqKE,1502
-mplang/backend/stablehlo.py,sha256=GOxy-qgOxyEdtBkt6LASKzaZnPewZhvHYSPOgFFXgIM,2612
+mplang/backend/stablehlo.py,sha256=RhKf6TUvjLrRvgtdVY2HxcRDGtjpKBobuBFMfsvZQOI,2937
 mplang/backend/tee.py,sha256=6kc7qTe8nWc3pr6iYtozEGLO8Umg-UBQLDiz6p3pdVg,1918
 mplang/core/__init__.py,sha256=lWxlEKfRwX7FNDzgyKZ1fiDMaCiqkyg0j5mKlZD_v7g,2244
 mplang/core/cluster.py,sha256=gqMJenvXUfHhE181Dd5JiUkD4nT07RLoicBnvsGmRkE,8598
@@ -41,7 +41,7 @@ mplang/frontend/base.py,sha256=rGtfBejcDh9mTRxOdJK5VUlG5vYiVJSir8X72X0Huvc,18264
 mplang/frontend/builtin.py,sha256=8qrlbe_SSy6QTXTnMG6_ADB8jSklVZGFBrkoR-p02FE,9368
 mplang/frontend/crypto.py,sha256=Nf8zT4Eko7MIs4R2tgZecKVd7d6Hvd_CGGmANhs3Ghs,3651
 mplang/frontend/ibis_cc.py,sha256=CTTbPPZ9hFnHuFDDIfgJHie1EdNnHmi5Ha1KsX0iYh8,4235
-mplang/frontend/jax_cc.py,sha256=ssP6rCvyWQ5VAr80-7z9QZUE2mWXyozJCGpq1dYQYY8,6374
+mplang/frontend/jax_cc.py,sha256=lMqaYD1tyM5DsStTNYifAXzhzsNM5nDiG3a61ygbWyc,7807
 mplang/frontend/phe.py,sha256=tDsCvStjVJ1Fs07yF3idkFnugUCA1zdFApPx7Uuulik,6795
 mplang/frontend/spu.py,sha256=7G6DaEfC5APSDhfeWSISTG_8tEcVbWth3XmjL8QUrVA,4994
 mplang/frontend/sql.py,sha256=DFdvjEPQX28VCRgUMeHYR0rwwOaoCH15bpvvlclLtHA,1999
@@ -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.dev148.dist-info/METADATA,sha256=3hkF4x8KZwg7PaGyfuAiDvp_qH9T1bhWS41rYWwl2Zs,16547
-mplang_nightly-0.1.dev148.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev148.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev148.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev148.dist-info/RECORD,,
+mplang_nightly-0.1.dev149.dist-info/METADATA,sha256=eZ_qGx1500gbGY5Ms8smHjQsyZ2Wxfa11z2mSeF0CSk,16547
+mplang_nightly-0.1.dev149.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev149.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev149.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev149.dist-info/RECORD,,