mplang-nightly 0.1.dev148__tar.gz → 0.1.dev149__tar.gz

{mplang_nightly-0.1.dev148 → mplang_nightly-0.1.dev149}/PKG-INFO

@@ -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 → mplang_nightly-0.1.dev149}/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_nightly-0.1.dev148 → mplang_nightly-0.1.dev149}/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 → mplang_nightly-0.1.dev149}/tests/frontend/test_jax_cc.py

@@ -244,21 +244,37 @@ class TestJax2StableHLO:
         assert cfunc.fn_text is not None
 
     def test_multiple_outputs(self):
-        """Test compilation with multiple outputs."""
+        """Test functions with multiple outputs."""
 
         def multi_output(x, y):
             return x + y, x - y, x * y
 
-        x = jnp.array([1.0, 2.0])
-        y = jnp.array([3.0, 4.0])
+        pfunc, out_tree = self._compile_with_transformer(
+            multi_output, jnp.array([1, 2]), jnp.array([3, 4])
+        )
 
-        cfunc, _out_tree = self._compile_with_transformer(multi_output, x, y)
+        assert len(pfunc.outs_info) == 3
+        assert out_tree is not None
 
-        assert len(cfunc.ins_info) == 2
-        assert len(cfunc.outs_info) == 3
+    def test_unused_parameter_elimination(self):
+        """Test that unused parameters are handled correctly via arg_keep_map."""
 
-        # All outputs should have the same shape as inputs
-        for out_info in cfunc.outs_info:
-            assert out_info.shape == x.shape
+        def func_with_unused(x, unused, z):
+            return x + z  # unused parameter eliminated by JAX
 
-        assert cfunc.fn_text is not None
+        x = jnp.array(1, dtype=jnp.int32)
+        unused = jnp.array(999, dtype=jnp.int32)
+        z = jnp.array(3, dtype=jnp.int32)
+
+        pfunc, _ = self._compile_with_transformer(func_with_unused, x, unused, z)
+
+        # Check that compilation succeeded
+        assert pfunc.fn_type == "mlir.stablehlo"
+        assert len(pfunc.ins_info) == 3  # Original input count
+
+        # If JAX eliminated unused parameters, arg_keep_map should be present
+        if "arg_keep_map" in pfunc.attrs:
+            keep_map = pfunc.attrs["arg_keep_map"]
+            assert isinstance(keep_map, list)
+            assert len(keep_map) < 3  # Should be fewer than original 3 params
+            assert 1 not in keep_map  # Index 1 (unused) should not be kept

mplang_nightly-0.1.dev149/tests/integration/test_unused_param_integration.py

@@ -0,0 +1,191 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Test unused parameter handling with mplang integration.
+This test verifies that functions with unused parameters work correctly
+after the arg_keep_map implementation.
+"""
+
+import jax.numpy as jnp
+
+import mplang
+import mplang.simp as simp
+
+
+def func_with_unused_params(a, unused_param, b, c):
+    """Function with unused parameter in the middle."""
+    return a + b + c
+
+
+def func_all_unused_returns_constant(a, unused1, unused2):
+    """Function where all parameters are unused - returns constant."""
+    return 42
+
+
+def func_first_last_unused(unused1, b, c, unused2):
+    """Function with unused parameters at start and end."""
+    return b * c
+
+
+class TestUnusedParameterHandling:
+    """Test suite for JAX unused parameter elimination handling."""
+
+    @staticmethod
+    def _extract_scalar(output):
+        """Extract scalar value from potentially wrapped output."""
+        if hasattr(output, "__iter__") and len(output) == 1:
+            output = output[0]
+        if hasattr(output, "item"):  # JAX array
+            output = output.item()
+        return output
+
+    def test_basic_unused_param(self):
+        """Test function with one unused parameter in middle position."""
+        sim = mplang.Simulator.simple(1)
+
+        # Create traced function
+        @mplang.function
+        def test_func():
+            # Test values - create inside traced context
+            a = simp.constant(1)
+            unused = simp.constant(999)  # This should be eliminated by JAX
+            b = simp.constant(2)
+            c = simp.constant(3)
+            return simp.run(func_with_unused_params)(a, unused, b, c)
+
+        expected = 6  # 1 + 2 + 3
+
+        # Compile and check that compilation succeeds
+        compiled = mplang.compile(sim, test_func)
+
+        # The function should compile successfully
+        assert compiled is not None
+
+        # Execute and verify result
+        result = mplang.evaluate(sim, test_func)
+        output = mplang.fetch(sim, result)
+
+        output = self._extract_scalar(output)
+
+        assert output == expected, f"Expected {expected}, got {output}"
+
+    def test_multiple_unused_params(self):
+        """Test function with multiple unused parameters."""
+        sim = mplang.Simulator.simple(1)
+
+        b_val = 5
+        c_val = 7
+        expected = b_val * c_val  # 35
+
+        @mplang.function
+        def test_func():
+            unused1 = simp.constant(100)
+            b = simp.constant(b_val)
+            c = simp.constant(c_val)
+            unused2 = simp.constant(200)
+            return simp.run(func_first_last_unused)(unused1, b, c, unused2)
+
+        result = mplang.evaluate(sim, test_func)
+        output = mplang.fetch(sim, result)
+        output = self._extract_scalar(output)
+
+        assert output == expected, f"Expected {expected}, got {output}"
+
+    def test_all_params_unused(self):
+        """Test function where all parameters are unused (returns constant)."""
+        sim = mplang.Simulator.simple(1)
+        expected = 42
+
+        @mplang.function
+        def test_func():
+            a = simp.constant(1)
+            unused1 = simp.constant(10)
+            unused2 = simp.constant(20)
+            return simp.run(func_all_unused_returns_constant)(a, unused1, unused2)
+
+        result = mplang.evaluate(sim, test_func)
+        output = mplang.fetch(sim, result)
+        output = self._extract_scalar(output)
+
+        assert output == expected, f"Expected {expected}, got {output}"
+
+    def test_no_unused_params(self):
+        """Test function with no unused parameters (regression test)."""
+        sim = mplang.Simulator.simple(1)
+
+        def func_all_used(a, b, c):
+            return a + b + c
+
+        @mplang.function
+        def test_func():
+            a = simp.constant(10)
+            b = simp.constant(20)
+            c = simp.constant(30)
+            return simp.run(func_all_used)(a, b, c)
+
+        result = mplang.evaluate(sim, test_func)
+        output = mplang.fetch(sim, result)
+        output = self._extract_scalar(output)
+
+        assert output == 60, f"Expected 60, got {output}"
+
+    def test_arg_keep_map_in_pfunc(self):
+        """Test that arg_keep_map is correctly stored in PFunction when needed."""
+        from mplang.frontend.jax_cc import jax2stablehlo
+
+        def func_with_unused(a, unused, b):
+            return a * b
+
+        # Create test inputs
+        a = jnp.array(2, dtype=jnp.int32)
+        unused = jnp.array(999, dtype=jnp.int32)
+        b = jnp.array(3, dtype=jnp.int32)
+
+        # Mock is_variable function
+        def is_variable(arg):
+            return True  # Treat all as variables for this test
+
+        # Call jax2stablehlo directly
+        pfunc, _, _ = jax2stablehlo(is_variable, func_with_unused, a, unused, b)
+
+        # Check that arg_keep_map is present when parameters are eliminated
+        if "arg_keep_map" in pfunc.attrs:
+            keep_map = pfunc.attrs["arg_keep_map"]
+            assert isinstance(keep_map, list)
+            assert len(keep_map) < 3  # Should be fewer than original 3 params
+            assert 1 not in keep_map  # Index 1 (unused) should not be in keep_map
+        else:
+            # If no elimination happened (possible with different JAX versions/optimizations)
+            pass
+
+    def test_different_dtypes_unused(self):
+        """Test unused parameter elimination with different data types."""
+        sim = mplang.Simulator.simple(1)
+
+        def func_mixed_types(int_used, float_unused, int_used2):
+            return int_used + int_used2  # float_unused is not used
+
+        @mplang.function
+        def test_func():
+            a = simp.constant(5)
+            unused_float = simp.constant(3.14)  # Different dtype, unused
+            c = simp.constant(7)
+            return simp.run(func_mixed_types)(a, unused_float, c)
+
+        result = mplang.evaluate(sim, test_func)
+        output = mplang.fetch(sim, result)
+        output = self._extract_scalar(output)
+
+        assert output == 12, f"Expected 12, got {output}"