feat(graph): introduce strip_pipeline_stage_metadata function and enhance graph signature handling

- Add `strip_pipeline_stage_metadata` to remove per-variable pipeline-stage hints from GraphDef.
- Update `_scan_graph_signature` to utilize the new function for cleaner graph signature processing.
- Refactor `_stack_module_states` to handle multiple graph definitions and ensure consistent signatures across modules.
- Enhance error handling for heterogeneous layers in `_stack_module_states`.
- Improve overall graph definition management in the core module.

Author: erfanzar

spectrax/core/containers.py

@@ -22,7 +22,7 @@
 import jax.numpy as jnp
 
 from ..sharding.mesh import current_mesh
-from .graph import GraphDef, ModuleNode, VarNode, iter_variables
+from .graph import GraphDef, ModuleNode, VarNode, iter_variables, strip_pipeline_stage_metadata
 from .module import Module, Opaque, _bump_graph_epoch, _graph_epoch
 from .paths import str_to_path
 from .registry import resolve_class
@@ -60,15 +60,16 @@ def _stack_module_states(items: list[Module], *, context: str) -> tuple[GraphDef
     if not items:
         raise ValueError(f"{context} requires at least one module")
     exports = [export(m) for m in items]
-    gdef = exports[0][0]
-    signature = _scan_graph_signature(gdef)
+    graph_defs = tuple(g for g, _state in exports)
+    signature = _scan_graph_signature(graph_defs[0])
     for index, (other_gdef, _state) in enumerate(exports[1:], start=1):
         if _scan_graph_signature(other_gdef) != signature:
             raise ValueError(
                 f"{context} requires every item to have the same graph structure; "
                 f"item 0 and item {index} differ. Use a Python loop for heterogeneous layers."
             )
     states = [s for _, s in exports]
+    gdef = _template_graphdef_without_mixed_stage_metadata(graph_defs)
     return gdef, jax.tree.map(lambda *vs: jnp.stack(vs, axis=0), *states)
 
 
@@ -125,26 +126,27 @@ def _scan_graph_signature(gdef: GraphDef) -> GraphDef:
     Returns:
         Return a graph signature suitable for repeated-layer scans.
     """
-    nodes = []
-    changed = False
-    for node in gdef.nodes:
-        if isinstance(node, VarNode):
-            metadata = tuple((k, v) for k, v in node.metadata if k != PIPELINE_STAGE_METADATA_KEY)
-            if metadata != node.metadata:
-                changed = True
-                node = VarNode(class_name=node.class_name, collection=node.collection, metadata=metadata)
-        nodes.append(node)
-    if not changed:
-        return gdef
-    return GraphDef(
-        nodes=tuple(nodes),
-        root=gdef.root,
-        var_refs=gdef.var_refs,
-        var_canonical=gdef.var_canonical,
-        shared_paths=gdef.shared_paths,
+    return strip_pipeline_stage_metadata(gdef)
+
+
+def _pipeline_stage_metadata_signature(gdef: GraphDef) -> tuple[tuple[tuple[str, object], ...], ...]:
+    """Return only the per-variable pipeline-stage metadata from ``gdef``."""
+    return tuple(
+        tuple((k, v) for k, v in node.metadata if k == PIPELINE_STAGE_METADATA_KEY)
+        for node in gdef.nodes
+        if isinstance(node, VarNode)
     )
 
 
+def _template_graphdef_without_mixed_stage_metadata(graph_defs: tuple[GraphDef, ...]) -> GraphDef:
+    """Use the first graph template, stripping stage metadata when it varies."""
+    template = graph_defs[0]
+    stage_signature = _pipeline_stage_metadata_signature(template)
+    if any(_pipeline_stage_metadata_signature(gdef) != stage_signature for gdef in graph_defs[1:]):
+        return strip_pipeline_stage_metadata(template)
+    return template
+
+
 def _scan_graph_topology_signature(gdef: GraphDef) -> tuple[object, ...]:
     """Return the state/child topology used to decide scan compatibility.
 
@@ -619,11 +621,11 @@ def _scan_static_template_signature(
     if len(graph_defs) == 1:
         return graph_defs[0]
     if family_keys is not None and all(key == family_keys[0] for key in family_keys[1:]):
-        return graph_defs[0]
+        return _template_graphdef_without_mixed_stage_metadata(graph_defs)
     key = _scan_graph_family_key(graph_defs[0])
     if any(_scan_graph_family_key(g) != key for g in graph_defs[1:]):
         return None
-    return graph_defs[0]
+    return _template_graphdef_without_mixed_stage_metadata(graph_defs)
 
 
 def _build_scan_plan_from_exports(exports: list[tuple[GraphDef, State]], *, context: str) -> _ScanPlan:

spectrax/core/graph.py

@@ -47,6 +47,7 @@
     "find",
     "iter_variables",
     "live_variables",
+    "strip_pipeline_stage_metadata",
     "tree_state",
     "update",
 ]
@@ -228,6 +229,35 @@ def canonical_path(self, ref_id: int) -> str:
         raise KeyError(f"No canonical path for ref_id {ref_id}")
 
 
+def strip_pipeline_stage_metadata(graphdef: GraphDef) -> GraphDef:
+    """Return ``graphdef`` with per-variable pipeline-stage hints removed.
+
+    A graph template can represent a stack or scan segment containing several
+    logical layers/stages. In that case one concrete layer's
+    ``pipeline_stage`` metadata is not a valid owner for the whole template.
+    """
+    from .stage_assignment import PIPELINE_STAGE_METADATA_KEY
+
+    nodes: list[Node] = []
+    changed = False
+    for node in graphdef.nodes:
+        if isinstance(node, VarNode):
+            metadata = tuple((k, v) for k, v in node.metadata if k != PIPELINE_STAGE_METADATA_KEY)
+            if metadata != node.metadata:
+                changed = True
+                node = VarNode(class_name=node.class_name, collection=node.collection, metadata=metadata)
+        nodes.append(node)
+    if not changed:
+        return graphdef
+    return GraphDef(
+        nodes=tuple(nodes),
+        root=graphdef.root,
+        var_refs=graphdef.var_refs,
+        var_canonical=graphdef.var_canonical,
+        shared_paths=graphdef.shared_paths,
+    )
+
+
 def _container_kind(m: Module) -> str:
     """Return the :attr:`ModuleNode.container_kind` string for ``m``.
 

spectrax/core/selector.py

@@ -505,8 +505,10 @@ def partition_state(self, module: Module, state: State) -> tuple[State, State]:
                 writer = state._writers.get((c, path))
                 if writer is not None:
                     if is_match:
+                        assert matched_writers is not None
                         matched_writers[(c, path)] = writer
                     else:
+                        assert rest_writers is not None
                         rest_writers[(c, path)] = writer
         return State._from_raw(matched_nested, writers=matched_writers), State._from_raw(
             rest_nested, writers=rest_writers

spectrax/runtime/mpmd/markers.py

@@ -41,6 +41,8 @@ def model(x):
 from jax.interpreters import ad, batching, mlir
 from jax.sharding import NamedSharding, PartitionSpec
 
+from spectrax._internal.logging import get_logger
+
 __all__ = [
     "cluster_jaxpr_by_markers",
     "has_stage_regions",
@@ -55,6 +57,9 @@ def model(x):
     "sxstage_region",
 ]
 
+logger = get_logger(__name__)
+_CLUSTER_PRUNE_DIAGNOSTICS = {"logged": 0}
+
 
 sxstage_iter_p = Primitive("sxstage_iter")
 sxstage_iter_p.multiple_results = True
@@ -900,6 +905,42 @@ def _collect_defined_vars_ordered(eqns: list[JaxprEqn]) -> list[Var]:
     return ordered
 
 
+def _prune_stage_jaxpr(sub: Jaxpr) -> Jaxpr:
+    """Drop stage equations that do not feed the stage outputs.
+
+    ``cluster_jaxpr_by_markers`` already computes a minimal-ish output tuple
+    for each stage boundary, but the stage body is later evaluated through
+    ``jax.core.eval_jaxpr``. Keeping dead equations in that private jaxpr can
+    force expensive auxiliary computations to survive tracing even when their
+    values are not part of the stage ABI. This pass performs plain reverse
+    liveness over the chosen outvars before the stage reaches ``jax.jit``.
+    """
+    needed: set[int] = {id(v) for v in sub.outvars if isinstance(v, Var)}
+    kept_rev: list[JaxprEqn] = []
+    for eqn in reversed(sub.eqns):
+        eqn_outvars = [v for v in eqn.outvars if isinstance(v, Var)]
+        effects = getattr(eqn, "effects", core.no_effects)
+        keep = bool(effects) or any(id(v) in needed for v in eqn_outvars)
+        if not keep:
+            continue
+        kept_rev.append(eqn)
+        for invar in eqn.invars:
+            if isinstance(invar, Var):
+                needed.add(id(invar))
+
+    pruned_eqns = list(reversed(kept_rev))
+    pruned_invars = [v for v in sub.invars if not isinstance(v, Var) or id(v) in needed]
+    if len(pruned_eqns) == len(sub.eqns) and len(pruned_invars) == len(sub.invars):
+        return sub
+    return Jaxpr(
+        constvars=list(sub.constvars),
+        invars=pruned_invars,
+        outvars=list(sub.outvars),
+        eqns=pruned_eqns,
+        effects=sub.effects,
+    )
+
+
 def _stage_region_spans(jaxpr: Jaxpr) -> tuple[tuple[int, int], ...]:
     """Return top-level equation spans covered by stage-region markers.
 
@@ -1309,6 +1350,8 @@ def collect_remat_eqns(
             defined_order_up_to_at[idx] = list(pre_order)
 
     clusters: list[Jaxpr] = []
+    dropped_eqns_total = 0
+    dropped_invars_total = 0
     for idx, (start, end) in enumerate(itertools.pairwise(boundaries)):
         base_eqns = [
             e for eqn_idx, e in enumerate(jaxpr.eqns[start:end], start=start) if eqn_idx not in boundary_marker_positions
@@ -1369,8 +1412,25 @@ def collect_remat_eqns(
             eqns=eqns,
             effects=core.no_effects,
         )
-        clusters.append(sub)
+        pruned = _prune_stage_jaxpr(sub)
+        dropped_eqns_total += len(sub.eqns) - len(pruned.eqns)
+        dropped_invars_total += len(sub.invars) - len(pruned.invars)
+        clusters.append(pruned)
         del idx
+    if dropped_eqns_total and _CLUSTER_PRUNE_DIAGNOSTICS.get("logged", 0) < 5:
+        try:
+            process_index = jax.process_index()
+        except Exception:
+            process_index = -1
+        if process_index == 0:
+            logger.warning(
+                "SpectraX MPMD marker clustering pruned %d dead stage equation(s) "
+                "and %d unused stage input(s) across %d stage(s).",
+                dropped_eqns_total,
+                dropped_invars_total,
+                len(clusters),
+            )
+            _CLUSTER_PRUNE_DIAGNOSTICS["logged"] = _CLUSTER_PRUNE_DIAGNOSTICS.get("logged", 0) + 1
     return clusters