onnx-ir 0.1.15__py3-none-any.whl

onnx_ir/passes/common/onnx_checker.py

@@ -0,0 +1,57 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Passes for debugging purposes."""
+
+from __future__ import annotations
+
+__all__ = [
+    "CheckerPass",
+]
+
+from typing import Literal
+
+import onnx  # noqa: TID251
+
+import onnx_ir as ir
+from onnx_ir.passes.common import _c_api_utils
+
+
+class CheckerPass(ir.passes.PassBase):
+    """Run onnx checker on the model."""
+
+    @property
+    def in_place(self) -> Literal[True]:
+        """This pass does not create a new model."""
+        return True
+
+    @property
+    def changes_input(self) -> Literal[False]:
+        """This pass does not change the input model."""
+        return False
+
+    def __init__(
+        self,
+        full_check: bool = False,
+        skip_opset_compatibility_check: bool = False,
+        check_custom_domain: bool = False,
+    ):
+        super().__init__()
+        self.full_check = full_check
+        self.skip_opset_compatibility_check = skip_opset_compatibility_check
+        self.check_custom_domain = check_custom_domain
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        """Run the onnx checker on the model."""
+
+        def _partial_check_model(proto: onnx.ModelProto) -> None:
+            """Partial function to check the model."""
+            onnx.checker.check_model(
+                proto,
+                full_check=self.full_check,
+                skip_opset_compatibility_check=self.skip_opset_compatibility_check,
+                check_custom_domain=self.check_custom_domain,
+            )
+
+        _c_api_utils.call_onnx_api(func=_partial_check_model, model=model)
+        # The model is not modified
+        return ir.passes.PassResult(model, False)

onnx_ir/passes/common/output_fix.py

@@ -0,0 +1,141 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Output fix pass for adding Identity nodes.
+
+- Graph inputs are directly used as outputs (without any intermediate nodes).
+- A value is used multiple times as a graph output (ensuring each output is unique).
+
+This ensures compliance with the ONNX specification for valid output configurations.
+"""
+
+from __future__ import annotations
+
+__all__ = [
+    "OutputFixPass",
+]
+
+import logging
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+class OutputFixPass(ir.passes.InPlacePass):
+    """Pass for adding Identity nodes to fix invalid output configurations.
+
+    This pass adds Identity nodes according to the following rules:
+
+    - If a graph input is directly used as a graph output (without any intermediate nodes),
+      insert an Identity node between them. The ONNX specification does not allow a graph
+      input to be directly used as a graph output without any processing nodes in between.
+    - If a value is used multiple times as graph outputs, insert Identity nodes for each
+      duplicate usage (keeping the first usage unchanged). This ensures each output value
+      is unique, as required by the ONNX specification.
+
+    This pass processes both the main graph and all subgraphs (e.g., in control flow operators).
+
+    Example transformations:
+        Direct input-to-output:
+            Before: input -> (direct connection) -> output
+            After:  input -> Identity -> output
+
+        Duplicate outputs:
+            Before: value -> [output1, output2]
+            After:  value -> output1, value -> Identity -> output2
+    """
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        """Main entry point for the output fix pass."""
+        modified = False
+
+        # Process the main graph
+        if _alias_multi_used_outputs(model.graph):
+            modified = True
+        if _alias_direct_outputs(model.graph):
+            modified = True
+
+        # Process functions
+        for function in model.functions.values():
+            if _alias_multi_used_outputs(function):
+                modified = True
+            if _alias_direct_outputs(function):
+                modified = True
+
+        return ir.passes.PassResult(model, modified=modified)
+
+
+def _alias_multi_used_outputs(graph_like: ir.Graph | ir.Function) -> bool:
+    """Insert Identity nodes for values that appear in the graph output list multiple times."""
+    modified = False
+
+    for graph in (graph_like, *graph_like.subgraphs()):
+        # Count usage of each output
+        seen: set[ir.Value] = set()
+
+        # Add Identity nodes for outputs used multiple times
+        for i, output in enumerate(graph.outputs):
+            if output not in seen:
+                # Skip the first occurrence
+                seen.add(output)
+                continue
+
+            # Create an Identity node
+            identity_node = ir.node("Identity", inputs=[output])
+            identity_output = identity_node.outputs[0]
+
+            # Copy metadata from the original output
+            # TODO: Use a better unique naming strategy if needed
+            identity_output.name = f"{output.name}_alias_{i}"
+            identity_output.shape = output.shape
+            identity_output.type = output.type
+            identity_output.metadata_props.update(output.metadata_props)
+            identity_output.doc_string = output.doc_string
+
+            # Add the node to the graph
+            graph.append(identity_node)
+            graph.outputs[i] = identity_output
+            logger.debug(
+                "Added Identity node for graph output '%s' used multiple times", output
+            )
+            modified = True
+    return modified
+
+
+def _alias_direct_outputs(graph_like: ir.Graph | ir.Function) -> bool:
+    """Insert Identity nodes for graph inputs used directly as outputs."""
+    modified = False
+
+    for graph in (graph_like, *graph_like.subgraphs()):
+        # Check each output to see if it's directly a graph input
+        outputs_to_fix: list[tuple[ir.Value, int]] = []
+        for i, output in enumerate(graph.outputs):
+            if output.is_graph_input():
+                outputs_to_fix.append((output, i))
+
+        # Add Identity nodes for each output that needs fixing
+        for output, index in outputs_to_fix:
+            # Create an Identity node
+            identity_node = ir.node("Identity", inputs=[output])
+            identity_output = identity_node.outputs[0]
+
+            # Copy metadata from the original output
+            # Preserve the original output name
+            identity_output.name = output.name
+            identity_output.shape = output.shape
+            identity_output.type = output.type
+            identity_output.metadata_props.update(output.metadata_props)
+            identity_output.doc_string = output.doc_string
+
+            # Create a new name for the old output
+            # TODO: Use a better unique naming strategy if needed
+            output.name = f"{output.name}_orig"
+
+            # Add the node to the graph
+            graph.append(identity_node)
+            graph.outputs[index] = identity_output
+
+            logger.debug("Added Identity node for graph input '%s' used as output", output)
+            modified = True
+
+    return modified

onnx_ir/passes/common/shape_inference.py

@@ -0,0 +1,112 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Shape inference pass using onnx.shape_inference."""
+
+from __future__ import annotations
+
+__all__ = [
+    "ShapeInferencePass",
+    "infer_shapes",
+]
+
+import logging
+
+import onnx  # noqa: TID251
+
+import onnx_ir as ir
+from onnx_ir.passes.common import _c_api_utils
+
+logger = logging.getLogger(__name__)
+
+
+def _merge_func(model: ir.Model, inferred_proto: onnx.ModelProto) -> bool:
+    """Merge the shape inferred model with the original model.
+
+    Args:
+        model: The original IR model.
+        inferred_proto: The ONNX model with shapes and types inferred.
+
+    Returns:
+        A tuple containing the modified model and a boolean indicating whether the model was modified.
+    """
+    inferred_model = ir.serde.deserialize_model(inferred_proto)
+    modified = False
+    for original_graph, inferred_graph in zip(model.graphs(), inferred_model.graphs()):
+        original_values = ir.convenience.create_value_mapping(original_graph)
+        inferred_values = ir.convenience.create_value_mapping(inferred_graph)
+        for name, value in original_values.items():
+            if name in inferred_values:
+                inferred_value = inferred_values[name]
+                if value.shape != inferred_value.shape and inferred_value.shape is not None:
+                    value.shape = inferred_value.shape
+                    modified = True
+                if value.dtype != inferred_value.dtype and inferred_value.dtype is not None:
+                    value.dtype = inferred_value.dtype
+                    modified = True
+            else:
+                logger.warning(
+                    "Value %s not found in inferred graph %s", name, inferred_graph.name
+                )
+    return modified
+
+
+class ShapeInferencePass(ir.passes.InPlacePass):
+    """This pass performs shape inference on the graph."""
+
+    def __init__(
+        self, check_type: bool = True, strict_mode: bool = True, data_prop: bool = True
+    ) -> None:
+        """Initialize the shape inference pass.
+
+        If inference fails, the model is left unchanged.
+
+        Args:
+            check_type: If True, check the types of the inputs and outputs.
+            strict_mode: If True, use strict mode for shape inference.
+            data_prop: If True, use data propagation for shape inference.
+        """
+        super().__init__()
+        self.check_type = check_type
+        self.strict_mode = strict_mode
+        self.data_prop = data_prop
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        def partial_infer_shapes(proto: onnx.ModelProto) -> onnx.ModelProto:
+            return onnx.shape_inference.infer_shapes(
+                proto,
+                check_type=self.check_type,
+                strict_mode=self.strict_mode,
+                data_prop=self.data_prop,
+            )
+
+        try:
+            inferred_model_proto = _c_api_utils.call_onnx_api(partial_infer_shapes, model)
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            logger.warning("Shape inference failed: %s. Model is left unchanged", exc_info=e)
+            return ir.passes.PassResult(model, False)
+
+        modified = _merge_func(model, inferred_model_proto)
+        return ir.passes.PassResult(model, modified=modified)
+
+
+def infer_shapes(
+    model: ir.Model,
+    *,
+    check_type: bool = True,
+    strict_mode: bool = True,
+    data_prop: bool = True,
+) -> ir.Model:
+    """Perform shape inference on the model.
+
+    Args:
+        model: The model to perform shape inference on.
+        check_type: If True, check the types of the inputs and outputs.
+        strict_mode: If True, use strict mode for shape inference.
+        data_prop: If True, use data propagation for shape inference.
+
+    Returns:
+        The model with shape inference applied.
+    """
+    return ShapeInferencePass(
+        check_type=check_type, strict_mode=strict_mode, data_prop=data_prop
+    )(model).model

onnx_ir/passes/common/topological_sort.py

@@ -0,0 +1,37 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Pass for topologically sorting the graphs."""
+
+from __future__ import annotations
+
+__all__ = [
+    "TopologicalSortPass",
+]
+
+
+import onnx_ir as ir
+
+
+class TopologicalSortPass(ir.passes.InPlacePass):
+    """Topologically sort graphs and functions in a model.
+
+    The sort is stable, preserving the relative order of nodes that are not
+    dependent on each other. Read more at :meth:`onnx_ir.Graph.sort`.
+    """
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        original_nodes = list(model.graph)
+        model.graph.sort()
+        sorted_nodes = list(model.graph)
+        for function in model.functions.values():
+            original_nodes.extend(function)
+            function.sort()
+            sorted_nodes.extend(function)
+
+        # Compare node orders to determine if any changes were made
+        modified = False
+        for node, new_node in zip(original_nodes, sorted_nodes):
+            if node is not new_node:
+                modified = True
+                break
+        return ir.passes.PassResult(model=model, modified=modified)

onnx_ir/passes/common/unused_removal.py

@@ -0,0 +1,215 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+__all__ = [
+    "RemoveUnusedNodesPass",
+    "RemoveUnusedFunctionsPass",
+    "RemoveUnusedOpsetsPass",
+]
+
+import logging
+
+import onnx  # noqa: TID251
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+def _remove_unused_optional_outputs(
+    node: ir.Node, graph_outputs: frozenset[ir.Value], onnx_opset_version: int
+) -> None:
+    try:
+        if node.domain not in {"", "onnx.ai"}:
+            return
+        op_schema = onnx.defs.get_schema(node.op_type, onnx_opset_version, domain=node.domain)
+    except Exception:  # pylint: disable=broad-exception-caught
+        logger.info(
+            "Failed to get schema for %s, skipping optional output removal",
+            node,
+            stack_info=True,
+        )
+        return
+
+    if node.op_type == "BatchNormalization":
+        # BatchNormalization op has 3 outputs: Y, running_mean, running_var
+        # If running_mean and running_var are not used, remove them, and the training_mode attribute
+        def is_used_output(i: int) -> bool:
+            if i < len(node.outputs):
+                val = node.outputs[i]
+                return val in graph_outputs or bool(val.uses())
+            return False
+
+        if is_used_output(1) or is_used_output(2):
+            return
+        if len(node.outputs) > 1:
+            node.outputs[1].name = ""
+        if len(node.outputs) > 2:
+            node.outputs[2].name = ""
+        node.attributes.pop("training_mode", None)
+        return
+
+    optional_info = []
+    for o in op_schema.outputs:
+        # Current ops do not have optional outputs if they have variable number of outputs
+        if o.option == onnx.defs.OpSchema.FormalParameterOption.Variadic:
+            return
+        optional_info.append(o.option == onnx.defs.OpSchema.FormalParameterOption.Optional)
+    # If no optional outputs in spec, skip delete operations
+    if len([o == 1 for o in optional_info]) == 0:
+        return
+
+    for i, out in enumerate(node.outputs):
+        if out not in graph_outputs and (not out.uses()) and optional_info[i] is True:
+            out.name = ""
+
+    # Remove trailing outputs with empty names by counting backwards
+    new_output_count = len(node.outputs)
+    for i in reversed(range(len(node.outputs))):
+        if not node.outputs[i].name:
+            new_output_count -= 1
+        else:
+            break
+    node.resize_outputs(new_output_count)
+
+
+def _remove_trailing_empty_inputs(node: ir.Node) -> None:
+    # Remove trailing None inputs
+    new_input_count = len(node.inputs)
+    for i in reversed(range(len(node.inputs))):
+        if node.inputs[i] is None:
+            new_input_count -= 1
+        else:
+            break
+    node.resize_inputs(new_input_count)
+
+
+def _remove_unused_nodes_in_graph_like(function_or_graph: ir.Function | ir.Graph) -> int:
+    graph_outputs = frozenset(function_or_graph.outputs)
+    onnx_opset_version = function_or_graph.opset_imports.get("", None)
+    count = 0
+    for node in reversed(function_or_graph):
+        removable = True
+        for output in node.outputs:
+            if output in graph_outputs or output.uses():
+                removable = False
+                break
+        if removable:
+            function_or_graph.remove(node, safe=True)
+            count += 1
+        else:
+            _remove_trailing_empty_inputs(node)
+            if onnx_opset_version is not None:
+                _remove_unused_optional_outputs(node, graph_outputs, onnx_opset_version)
+            for attr in node.attributes.values():
+                if attr.type == ir.AttributeType.GRAPH:
+                    count += _remove_unused_nodes_in_graph_like(attr.as_graph())
+                elif attr.type == ir.AttributeType.GRAPHS:
+                    for graph in attr.as_graphs():
+                        count += _remove_unused_nodes_in_graph_like(graph)
+    return count
+
+
+class RemoveUnusedNodesPass(ir.passes.InPlacePass):
+    """Pass for removing unused nodes and initializers (dead code elimination).
+
+    This pass does not modify the model signature (inputs and outputs). It ensures
+    that unused nodes and initializers are removed while preserving the original
+    contract of the model.
+    """
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        count = _remove_unused_nodes_in_graph_like(model.graph)
+        graph_outputs = frozenset(model.graph.outputs)
+        graph_inputs = frozenset(model.graph.inputs)
+        initializers = model.graph.initializers
+        for init in list(initializers.values()):
+            if not (init.uses() or init in graph_outputs or init in graph_inputs):
+                assert init.name is not None
+                del initializers[init.name]
+                count += 1
+        for function in model.functions.values():
+            count += _remove_unused_nodes_in_graph_like(function)
+        if count:
+            logger.info("Removed %s unused nodes", count)
+        return ir.passes.PassResult(model, modified=bool(count))
+
+
+class RemoveUnusedFunctionsPass(ir.passes.InPlacePass):
+    def __init__(self):
+        super().__init__()
+        self._used: set[ir.OperatorIdentifier] | None = None
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        self._used = set()
+        for node in ir.traversal.RecursiveGraphIterator(model.graph):
+            self._call_node(model, node)
+
+        # Update the model to remove unused functions
+        unused = set(model.functions) - self._used
+        if not unused:
+            logger.info("No unused functions to remove")
+            return ir.passes.PassResult(model, modified=False)
+
+        for op_identifier in unused:
+            del model.functions[op_identifier]
+
+        logger.info("Removed %s unused functions", len(unused))
+        logger.debug("Functions left: %s", list(model.functions))
+        logger.debug("Functions removed: %s", unused)
+
+        self._used = None
+        return ir.passes.PassResult(model, modified=bool(unused))
+
+    def _call_function(self, model: ir.Model, function: ir.Function) -> None:
+        assert self._used is not None
+        if function.identifier() in self._used:
+            # The function and its nodes are already recorded as used
+            return
+        self._used.add(function.identifier())
+        for node in ir.traversal.RecursiveGraphIterator(function):
+            self._call_node(model, node)
+
+    def _call_node(self, model: ir.Model, node: ir.Node) -> None:
+        op_identifier = node.op_identifier()
+        if op_identifier not in model.functions:
+            return
+        self._call_function(model, model.functions[op_identifier])
+
+
+class RemoveUnusedOpsetsPass(ir.passes.InPlacePass):
+    """Remove unused opset imports from the model and functions.
+
+    Attributes:
+        process_functions: Whether to process functions in the model. If True, the pass will
+            remove unused opset imports from functions as well. If False, only the main graph
+            will be processed.
+    """
+
+    def __init__(self, process_functions: bool = True):
+        super().__init__()
+        self.process_functions = process_functions
+
+    def _process_graph_like(
+        self, graph_like: ir.Graph | ir.Function, used_domains: set[str]
+    ) -> bool:
+        for node in ir.traversal.RecursiveGraphIterator(graph_like):
+            used_domains.add(node.domain)
+        unused = set(graph_like.opset_imports) - used_domains
+        for domain in unused:
+            del graph_like.opset_imports[domain]
+        return bool(unused)
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        # Record domains of all functions
+        used_domains = {""}  # By default always retain the onnx (default) domain
+        for function in model.functions.values():
+            used_domains.add(function.domain)
+        modified = self._process_graph_like(model.graph, used_domains=used_domains)
+
+        if self.process_functions:
+            for function in model.functions.values():
+                modified |= self._process_graph_like(function, used_domains={""})
+
+        return ir.passes.PassResult(model, modified=modified)

onnx_ir/py.typed

@@ -0,0 +1 @@
+