graphutils API Reference

graphutils

Class: _LazyModel

Proxy that defers the expensive _make_graph() call until the result is actually used. When callers do _, model = replacer.replace_subgraph_by_edges(...) and immediately discard model, the rebuild never happens.

Note: resolution captures the DAG state at resolution time, not creation time. Only the last returned model from a sequence of replace_* calls should be resolved; intermediate models will reflect the final DAG state if resolved after further mutations.

Methods Overview

• __init__(builder) -
• __getattr__(name) -
• __repr__() -
• __iter__() -
• __bool__() -

Methods Details

__init__(builder)

__getattr__(name)

__repr__()

__iter__()

__bool__()

Class: MatchInfo

Inherits from: dict

Dictionary-like class to store information about matched patterns in an ONNX graph.

MatchInfo extends Python's dict to provide specialized functionality for pattern matching in ONNX graphs. It stores mappings from FilterNode objects to the corresponding matched NodeProto objects, and provides additional methods for accessing matched nodes by name and filtering results.

The primary purpose is to store the results of pattern matching operations and provide convenient ways to access and manipulate those results, especially allowing lookup by FilterNode name as well as by FilterNode object.

The user is not expected to directly operate on any objects of this class.

Methods Overview

• __init__() -
• __contains__(key) -
• __getitem__(key) -
• all_matched_nodes(exclude_wildcards) -
• by_filter_name() -

Methods Details

__init__(*args, **kwargs)

__contains__(key)

__getitem__(key)

all_matched_nodes(exclude_wildcards)

by_filter_name()

Class: FilterWildcard

A filter that matches any node in an ONNX graph.

This class is used in pattern matching to create wildcards that can match any node in the graph, allowing for flexible pattern definitions. When used within a FilterNode's operands list, it indicates that any node can match at that position in the pattern.

Methods Overview

• __init__(name) - Initialize a wildcard filter for pattern matching.
• match(node, onnx_model) - Match any node in the ONNX graph.

Methods Details

__init__(name)

Initialize a wildcard filter for pattern matching.

Parameters:

• name : str A name for the wildcard, used to retrieve the matched node from the MatchInfo. This allows referencing the matched node by name when processing matches.

match(node, onnx_model: OnnxModel)

Match any node in the ONNX graph.

This method always returns a successful match, associating this wildcard with the provided node in the returned MatchInfo.

Parameters:

• node : Union[onnx.NodeProto, onnx.TensorProto] The node to match. This can be any node in the ONNX graph.
• onnx_model : OnnxModel The ONNX model containing the node. Not used by this matcher.

Returns:

• MatchInfo A MatchInfo object with this wildcard mapped to the provided node.

Examples:

>>> wildcard = FilterWildcard("X")
>>> match_info = wildcard.match(some_node, onnx_model)
>>> # match_info now contains {wildcard: some_node}
>>> # The node can be retrieved using either the wildcard object or its name
>>> matched_node = match_info[wildcard]  # Using the wildcard object
>>> matched_node = match_info["X"]       # Using the wildcard name

Class: FilterNode

A class representing a filter node that is used for pattern matching in an ONNX graph.

This class allows creating patterns that can match specific structures in the ONNX graph and provides functionality for pattern matching with attribute checking, recursive operand checking, and support for commutative operations.

Methods Overview

• __init__(op_type, operands, attr_setting, keep_const_inputs, process_const_callback, process_attr_callback, name, commutative) - Initialize a new instance of the FilterNode class.
• match(node, onnx_model) - Match this filter against a node in the ONNX graph.
• __or__(other) - Create a union of this filter node with another, matching if either matches.

Methods Details

__init__(op_type: str, operands: list, attr_setting: int, keep_const_inputs: list[int], process_const_callback: Callable[[dict, int], dict], process_attr_callback: Callable[[dict, int], dict], name, commutative, **kwargs)

Initialize a new instance of the FilterNode class.

Parameters:

• op_type : str The operator type to match (e.g., "Conv", "Relu")
• operands : list The list of input FilterNodes to this filter node. These define the pattern structure. Use "*" or FilterWildcard to match any node.
• attr_setting : int Controls attribute matching behavior: -1 (Setting.MIN): Match if node has at least the filter's attributes 0 (Setting.EXACT): Match only if attributes are exactly the same 1 (Setting.MAX): Match if filter has at least the node's attributes Default is Setting.MIN (-1).
• keep_const_inputs : list[int] Indices of constant inputs that should be kept when replacing the matched pattern with a custom operator. If None, all constant inputs will be kept.
• process_const_callback : Callable[[dict, int], dict] Callback function to process constant inputs when replacing with a custom operator. The first argument is a dict mapping input indices to (name, value) tuples. The second argument is the counter of the custom op being created.
• process_attr_callback : Callable[[dict, int], dict] Callback function to process attributes as inputs when creating custom operators. Similar to process_const_callback but for attributes.
• name : str, optional Name to identify this filter node, used to retrieve matched nodes from the MatchInfo.
• commutative : bool, optional Whether the operation is commutative. If True, both orders of operands will be attempted during matching. Only applicable for binary operations. **kwargs Attributes to be matched with the ONNX node. These are specified as key-value pairs where the key is the attribute name and the value is the expected attribute value.

match(node, onnx_model: OnnxModel)

Match this filter against a node in the ONNX graph.

This recursively matches the filter pattern against a node and its operands.

Parameters:

• node : onnx.NodeProto The node to match against.
• onnx_model : OnnxModel The ONNX model containing the node.

Returns:

• MatchInfo Information about the match, or an empty MatchInfo if no match was found. MatchInfo maps filter nodes to their corresponding matched ONNX nodes.

Notes:

The matching process works as follows:

1. Check if the node's op_type matches this filter's op_type
2. Check if the number of inputs matches
3. Check if the node's attributes match according to the attr_setting
4. Recursively check if all operands match
5. For commutative operations, try matching both operand orders

Examples:

>>> conv = FilterNode("Conv", [FilterWildcard("X")], -1)
>>> match_info = conv.match(some_node, onnx_model)
>>> if match_info:
>>>     print(f"Found Conv node: {match_info[conv].name}")
>>>     print(f"With input: {match_info['X'].name}")

__or__(other)

Create a union of this filter node with another, matching if either matches.

Parameters:

• other : FilterNode Another filter node to union with this one.

Returns:

• FilterNodeUnion A new filter that matches if either this filter or the other filter matches.

Examples:

>>> # Match either Conv or Pool operations
>>> conv = FilterNode("Conv", [FilterWildcard("X")], -1)
>>> pool = FilterNode("MaxPool", [FilterWildcard("X")], -1)
>>> conv_or_pool = conv | pool
>>>
>>> # Use in matching
>>> model.match(conv_or_pool, callback_function)

Class: CustomOpReplacer

Class for replacing patterns in an ONNX graph with custom operators.

This class provides methods for finding and replacing patterns in an ONNX graph with custom operators, allowing for code specialization and optimization.

Methods Overview

• __init__(model) - Initialize a CustomOpReplacer with an ONNX model.
• make_graph() - Serialize the current DAG state to an ONNX ModelProto.
• replace_subgraph(output_nodes, input_nodes, ccl_func_name, element_wise, keep_constants, dimension_value, ccl_func_template_params, fixed_point_frac_bits, process_const_callback, match_callback, allow_io_in_l2_mem, reserved_ocm, reserved_ext, persistent_ext, needs_iter_var, op_name_base) - Replace a subgraph with a custom operator.
• replace_subgraph_by_edges(output_edges, input_edges, ccl_func_name, element_wise, keep_constants, dimension_value, ccl_func_template_params, fixed_point_frac_bits, process_const_callback, match_callback, allow_io_in_l2_mem, reserved_ocm, reserved_ext, persistent_ext, needs_iter_var, op_name_base) - Replace a subgraph defined by input and output edges with a custom operator.
• extract_subgraph_by_name_matching(node_name_patterns, subgraph_name) - Extract a subgraph containing nodes whose names match specified patterns.
• replace_by_name_matching(node_name_patterns, ccl_func_name, element_wise, keep_constants, dimension_value, ccl_func_template_params, fixed_point_frac_bits, process_const_callback, match_callback, allow_io_in_l2_mem, reserved_ocm, reserved_ext, persistent_ext, needs_iter_var, op_name_base) - Replace nodes matching name patterns with a custom operator.
• replace_all(filter_leaf, ccl_func_name, element_wise, keep_constants, dimension_value, ccl_func_template_params, fixed_point_frac_bits, process_const_callback, match_callback, allow_io_in_l2_mem, reserved_ocm, reserved_ext, persistent_ext, needs_iter_var, op_name_base) - Replace all occurrences of a pattern in the model with custom operators.
• fold_quant_nodes_into_custom_ops() - Fold quantization nodes into custom operators.

Methods Details

__init__(model)

Initialize a CustomOpReplacer with an ONNX model.

Parameters:

• model : onnx.ModelProto The ONNX model to operate on.

make_graph() -> onnx.ModelProto

Serialize the current DAG state to an ONNX ModelProto.

Returns:

• onnx.ModelProto

replace_subgraph(output_nodes: list[Union[str, onnx.NodeProto]], input_nodes: list[str], ccl_func_name: str, element_wise: bool, keep_constants: list[str], dimension_value: dict[str, int], ccl_func_template_params: list, fixed_point_frac_bits: int, process_const_callback: Callable[[dict], dict], match_callback: Callable[list[DagNode], dict], allow_io_in_l2_mem: bool, reserved_ocm: int, reserved_ext: int, persistent_ext: int, needs_iter_var: bool, op_name_base: str) -> onnx.ModelProto

Replace a subgraph with a custom operator.

This method identifies a subgraph defined by input and output nodes and replaces it with a custom operator. The subgraph is extracted, converted to a custom operation implementation, and replaced in the original graph with a QuadricCustomOp node.

Parameters:

• output_nodes : list[Union[str, onnx.NodeProto]] One or more node names or NodeProto objects that define the outputs of the subgraph. These are boundary nodes whose outputs will become the outputs of the custom operator.
• input_nodes : list[str] Node names that define the inputs of the subgraph. These are boundary nodes whose inputs will become the inputs of the custom operator.
• ccl_func_name : str Name of the custom function to be implemented in CCL (Chimera Compute Language). This name will be used in the generated custom operator.
• element_wise : bool Whether the operation is elementwise. This affects optimization decisions.
• keep_constants : list[str] Names of constants to keep in the custom op. If None, all constants are kept.
• dimension_value : dict[str, int] Map from dimension parameter names to static dimension values. Used to resolve dynamic dimensions to static values.
• ccl_func_template_params : list Template parameters for the custom function. Default is [""].
• fixed_point_frac_bits : int Number of fractional bits for fixed-point types.
• process_const_callback : Callable[[dict], dict] Callback function to process constant inputs when creating the custom operator. The callback receives a dictionary of constant inputs and can modify them.
• match_callback : Callable[list[DagNode], dict] Callback function called after a match to modify attributes of the custom operator. The callback receives the list of matched nodes and can return additional attributes.
• allow_io_in_l2_mem : bool Whether to allow inputs/outputs in L2 memory instead of external memory.
• reserved_ocm : int Reserved on-chip memory (OCM) size in bytes.
• reserved_ext : int Reserved ext/ddr memory size in bytes.
• persistent_ext : int Reserved persistent ext/ddr memory size in bytes.
• needs_iter_var : bool Whether the custom op should recieve the iteration count variable.
• op_name_base : str Base name for the custom operator. If None, ccl_func_name is used.

Returns:

• onnx.ModelProto A tuple containing (subgraph, modified_model), where:
  • subgraph: The extracted subgraph as an ONNX model
  • modified_model: The original model with the subgraph replaced by the custom operator

Raises:

ValueError If a subgraph input node is a primary graph input, indicating incorrect boundary definition. If specified input/output nodes are not found in the graph.

Notes:

• The subgraph must be properly defined with input and output boundary nodes.
• Constant nodes within the subgraph will be extracted and included in the custom operator.
• The method uses a breadth-first search starting from output nodes to identify all nodes in the subgraph, stopping at input nodes.
• For handling complex subgraphs with shared constants, consider using replace_subgraph_by_edges which defines boundaries by edge names instead of node names.

Examples:

Replace a Conv->Relu pattern with a custom operator:

>>> replacer = CustomOpReplacer(model)
>>> subgraph, modified = replacer.replace_subgraph(
>>>     ["relu_output_node"],  # Output node
>>>     ["conv_node"],         # Input node
>>>     "ConvRelu",            # Custom function name
>>>     False                  # Not element-wise
>>> )

Replace a subgraph with multiple inputs and outputs:

>>> subgraph, modified = replacer.replace_subgraph(
>>>     ["output1", "output2"],     # Multiple output nodes
>>>     ["input1", "input2"],       # Multiple input nodes
>>>     "CustomOp",
>>>     False,
>>>     keep_constants=["weight", "bias"]  # Keep these constants in the custom op
>>> )

Replace a subgraph with dimension resolution for dynamic shapes:

>>> subgraph, modified = replacer.replace_subgraph(
>>>     ["output_node"],
>>>     ["input_node"],
>>>     "CustomOp",
>>>     True,  # Element-wise
>>>     dimension_value={"NonMaxSuppression_257_o0__d0": 64}  # Resolve dynamic dimension
>>> )

replace_subgraph_by_edges(output_edges: list[str], input_edges: list[str], ccl_func_name: str, element_wise: bool, keep_constants: list[str], dimension_value: dict[str, int], ccl_func_template_params: list, fixed_point_frac_bits: int, process_const_callback: Callable[[dict], dict], match_callback: Callable[list[DagNode], dict], allow_io_in_l2_mem: bool, reserved_ocm: int, reserved_ext: int, persistent_ext: int, needs_iter_var: bool, op_name_base: str) -> onnx.ModelProto

Replace a subgraph defined by input and output edges with a custom operator.

This method identifies a subgraph defined by its boundary edges (rather than nodes) and replaces it with a custom operator. This approach is useful when working with complex graphs where boundary nodes may have multiple inputs/outputs.

Parameters:

• output_edges : list[str] One or more edge names that define the outputs of the subgraph. These are tensor names that will become outputs of the custom operator.
• input_edges : list[str] Edge names that define the inputs of the subgraph. These are tensor names that will become inputs of the custom operator.
• ccl_func_name : str Name of the custom function to be implemented in CCL (Chimera Compute Language). This name will be used in the generated custom operator.
• element_wise : bool Whether the operation is elementwise. This affects optimization decisions.
• keep_constants : list[str] Names of constants to keep in the custom op. If None, all constants are kept.
• dimension_value : dict[str, int] Map from dimension parameter names to static dimension values. Used to resolve dynamic dimensions to static values.
• ccl_func_template_params : list Template parameters for the custom function. Default is [""].
• fixed_point_frac_bits : int Number of fractional bits for fixed-point types.
• process_const_callback : Callable[[dict], dict] Callback function to process constant inputs when creating the custom operator. The callback receives a dictionary of constant inputs and can modify them.
• match_callback : Callable[list[DagNode], dict] Callback function called after a match to modify attributes of the custom operator. The callback receives the list of matched nodes and can return additional attributes.
• allow_io_in_l2_mem : bool Whether to allow inputs/outputs in L2 memory instead of external memory.
• reserved_ocm : int Reserved on-chip memory (OCM) size in bytes.
• reserved_ext : int Reserved ext/ddr memory size in bytes.
• persistent_ext : int Reserved persistent ext/ddr memory size in bytes.
• needs_iter_var : bool Whether the custom op should recieve the iteration count variable.
• op_name_base : str Base name for the custom operator. If None, ccl_func_name is used.

Returns:

• onnx.ModelProto A tuple containing (subgraph, modified_model), where:
  • subgraph: The extracted subgraph as an ONNX model
  • modified_model: The original model with the subgraph replaced by the custom operator

Raises:

ValueError If specified edge names are not found in the graph.

Notes:

• This method is particularly useful for subgraphs where the boundary is more naturally defined by tensor edges rather than nodes.
• The method uses a breadth-first search starting from output edges and working backward through the graph, stopping at input edges.
• More precise than replace_subgraph in some cases, especially when nodes have multiple inputs or outputs.

Examples:

Replace a subgraph defined by input/output edges:

>>> replacer = CustomOpReplacer(model)
>>> subgraph, modified = replacer.replace_subgraph_by_edges(
>>>     ["/output/edge"],  # Output edge
>>>     ["/input/edge"],   # Input edge
>>>     "CustomOp",
>>>     False
>>> )

Replace a complex subgraph with multiple input and output edges:

>>> subgraph, modified = replacer.replace_subgraph_by_edges(
>>>     ["/layers.0/attention_norm/Mul_output_0"],
>>>     ["/tok_embeddings/Gather_output_0"],
>>>     "customOpFuctionName",
>>>     False
>>> )

extract_subgraph_by_name_matching(node_name_patterns: Union[str, list[str]], subgraph_name)

Extract a subgraph containing nodes whose names match specified patterns.

This method identifies all nodes in the graph whose names match the provided regex patterns and extracts them as a subgraph. The extracted subgraph includes all matched nodes and the edges between them, with proper handling of input and output boundaries.

Parameters:

• node_name_patterns : Union[str, list[str]] One or more regex patterns to match node names. Any node whose name matches any of these patterns will be included in the extracted subgraph.
• subgraph_name : str, optional Name to assign to the extracted subgraph. If empty, defaults to the current ccl_func_name set in the CustomOpReplacer.

Returns:

• Tuple[onnx.ModelProto, OrderedSet[DagNode]] A tuple containing:
  • The extracted subgraph as an ONNX model
  • An OrderedSet of DagNode objects that were matched and included in the subgraph

Notes:

• This is a lower-level method primarily used by replace_by_name_matching but can be used directly when you want to extract a subgraph without replacing it.
• The regex patterns are wrapped with parentheses, so your pattern should not include capturing groups unless you intend them to be nested.
• The extracted subgraph includes all necessary input/output ValueInfoProtos and initializers needed by the matched nodes.
• If the matched nodes don't form a connected subgraph, the result will include all components with proper input/output boundaries.

Examples:

Extract nodes matching a single pattern:

>>> replacer = CustomOpReplacer(model)
>>> subgraph, matched_nodes = replacer.extract_subgraph_by_name_matching(
>>>     "TFNodes/.*",      # Regex pattern for TensorFlow nodes
>>>     "tf_subgraph"      # Name for the extracted subgraph
>>> )

Extract nodes matching multiple patterns:

>>> subgraph, matched_nodes = replacer.extract_subgraph_by_name_matching(
>>>     ["layer1_.*", "layer2_.*"],  # Multiple patterns
>>>     "combined_layers"            # Name for the subgraph
>>> )

replace_by_name_matching(node_name_patterns: Union[str, list[str]], ccl_func_name: str, element_wise: bool, keep_constants: list[str], dimension_value: dict[str, int], ccl_func_template_params: list, fixed_point_frac_bits: int, process_const_callback: Callable[[dict, int], dict], match_callback: Callable[list[DagNode], dict], allow_io_in_l2_mem: bool, reserved_ocm: int, reserved_ext: int, persistent_ext: int, needs_iter_var: bool, op_name_base: str) -> onnx.ModelProto

Replace nodes matching name patterns with a custom operator.

This method identifies nodes in the graph whose names match specified regex patterns and replaces them with a single custom operator. This is useful for targeting specific nodes based on naming conventions without having to define explicit pattern structures.

Parameters:

• node_name_patterns : Union[str, list[str]] One or more regex patterns to match node names. Any node whose name matches any of these patterns will be included in the subgraph to be replaced.
• ccl_func_name : str Name of the custom function to be implemented in CCL (Chimera Compute Language). This name will be used in the generated custom operator.
• element_wise : bool Whether the operation is elementwise. This affects optimization decisions.
• keep_constants : list[str] Names of constants to keep in the custom op. If None, all constants are kept.
• dimension_value : dict[str, int] Map from dimension parameter names to static dimension values. Used to resolve dynamic dimensions to static values.
• ccl_func_template_params : list Template parameters for the custom function. Default is [""].
• fixed_point_frac_bits : int Number of fractional bits for fixed-point types.
• process_const_callback : Callable[[dict, int], dict] Callback function to process constant inputs when creating the custom operator. The callback receives a dictionary of constant inputs and can modify them. The second parameter is the counter of the custom op being created.
• match_callback : Callable[list[DagNode], dict] Callback function called after a match to modify attributes of the custom operator. The callback receives the list of matched nodes and can return additional attributes.
• allow_io_in_l2_mem : bool Whether to allow inputs/outputs in L2 memory instead of external memory.
• reserved_ocm : int Reserved on-chip memory (OCM) size in bytes.
• reserved_ext : int Reserved ext/ddr memory size in bytes.
• persistent_ext : int Reserved persistent ext/ddr memory size in bytes.
• needs_iter_var : bool Whether the custom op should recieve the iteration count variable.
• op_name_base : str Base name for the custom operator. If None, ccl_func_name is used.

Returns:

• onnx.ModelProto A tuple containing (subgraph, modified_model), where:
  • subgraph: The extracted subgraph as an ONNX model
  • modified_model: The original model with the subgraph replaced by the custom operator

Notes:

• The regex patterns are wrapped with parentheses, so your pattern should not include capturing groups unless you intend them to be nested.
• All matched nodes and their interconnections form the subgraph to be replaced.
• The subgraph is automatically determined based on the dependency structure of the matched nodes.
• This method is particularly useful for targeting specific model sections where the node naming follows a convention, such as model layers with specific prefixes.
• For more control over the exact subgraph structure, consider using replace_subgraph or replace_subgraph_by_edges.

Examples:

Replace all nodes with names matching a pattern:

>>> replacer = CustomOpReplacer(model)
>>> subgraph, modified = replacer.replace_by_name_matching(
>>>     "TFNodes/.*",  # Regex pattern to match node names
>>>     "custom_op",   # Custom function name
>>>     False          # Not element-wise
>>> )

Replace nodes with multiple patterns:

>>> subgraph, modified = replacer.replace_by_name_matching(
>>>     ["layer1_.*", "layer2_.*"],  # Multiple regex patterns
>>>     "CustomLayers",
>>>     False
>>> )

Replace nodes with specific configuration:

>>> subgraph, modified = replacer.replace_by_name_matching(
>>>     "/layers.0/attention_norm.*",
>>>     "custom_op",
>>>     False,
>>>     keep_constants=["weight", "bias"],
>>>     ccl_func_template_params=["int8", "fp16"],
>>>     fixed_point_frac_bits=8
>>> )

replace_all(filter_leaf: FilterNode, ccl_func_name: str, element_wise: bool, keep_constants: list[str], dimension_value: dict[str, int], ccl_func_template_params: list, fixed_point_frac_bits: int, process_const_callback: Callable[[dict, int], dict], match_callback: Callable[[list[DagNode]], dict], allow_io_in_l2_mem: bool, reserved_ocm: int, reserved_ext: int, persistent_ext: int, needs_iter_var: bool, op_name_base: str) -> onnx.ModelProto

Replace all occurrences of a pattern in the model with custom operators.

This method traverses the entire graph, finds all occurrences of a specified pattern, and replaces each occurrence with a custom operator. The pattern is defined as a tree of FilterNode objects that describe the operations and their connections.

Parameters:

• filter_leaf : FilterNode The root of the filter pattern to match. This is typically the output node of the pattern, with child FilterNodes representing inputs. The method will search for all occurrences of this pattern in the graph.
• ccl_func_name : str Name of the custom function to be implemented in CCL (Chimera Compute Language). This name will be used in the generated custom operators.
• element_wise : bool Whether the operation is elementwise. This affects optimization decisions.
• keep_constants : list[str] Names of constants to keep in the custom op. If None, all constants are kept.
• dimension_value : dict[str, int] Map from dimension parameter names to static dimension values. Used to resolve dynamic dimensions to static values.
• ccl_func_template_params : list Template parameters for the custom function. Default is [""].
• fixed_point_frac_bits : int Number of fractional bits for fixed-point types.
• process_const_callback : Callable[[dict, int], dict] Callback function to process constant inputs when creating custom operators. The callback receives a dictionary mapping input indices to (name, value) tuples, and returns a similar dictionary. The second parameter is the counter of the custom op being created.
• match_callback : Callable[[list[DagNode]], dict] Callback function called after each match to modify attributes of the custom operator. The callback receives the list of matched nodes and can return additional attributes.
• allow_io_in_l2_mem : bool Whether to allow inputs/outputs in L2 memory instead of external memory.
• reserved_ocm : int Reserved on-chip memory (OCM) size in bytes.
• reserved_ext : int Reserved ext/ddr memory size in bytes.
• persistent_ext : int Reserved persistent ext/ddr memory size in bytes.
• needs_iter_var : bool Whether the custom op should recieve the iteration count variable.
• op_name_base : str Base name for the custom operators. If None, ccl_func_name is used. Each operator will have a unique suffix appended to this base name.

Returns:

• onnx.ModelProto The modified ONNX model with all occurrences of the pattern replaced by custom operators.

Notes:

• The matching process is recursive and will match nested patterns.
• Commutative operations (like Add) can be matched with both operand orders if the FilterNode has commutative=True.
• If a match is found, all nodes in the matched pattern will be replaced with a single custom operator node.
• Each replacement creates a subgraph for the custom operator implementation.
• For each match, constant inputs are collected and processed according to the keep_constants and process_const_callback parameters.
• If a node in the pattern has uses outside the pattern, the match will be rejected unless those external uses are explicitly handled.
• The method processes the graph in a topological order to ensure that all pattern instances are found.

Examples:

Replace all Conv->Relu patterns with custom operators:

>>> # Define the pattern
>>> inp = FilterWildcard("inp")
>>> conv = FilterNode("Conv", [inp], -1)
>>> relu = FilterNode("Relu", [conv], -1)
>>>
>>> # Replace all occurrences
>>> replacer = CustomOpReplacer(model)
>>> modified_model = replacer.replace_all(
>>>     relu,            # Leaf node of the pattern
>>>     "ConvRelu",      # Custom function name
>>>     False            # Not element-wise
>>> )

Replace all Conv nodes with specific padding:

>>> # Define a specific Conv pattern with pads=[0,0,0,0]
>>> filter_leaf = FilterNode("Conv", [FilterWildcard("X")], -1, pads=[0, 0, 0, 0])
>>>
>>> # Replace all matching Convs
>>> modified_model = replacer.replace_all(
>>>     filter_leaf,
>>>     "CustomConv",
>>>     False,
>>>     keep_constants=["weight", "bias"]
>>> )

Replace patterns with custom attribute processing:

>>> # Define pattern
>>> filter_leaf = FilterNode("QLinearConv", [FilterWildcard("X")], -1, dilations=[3, 3])
>>>
>>> # Define callback for template parameters
>>> def process_attributes(nodes_dict):
>>>     return {"custom_attr": "value"}
>>>
>>> # Replace with template parameters
>>> modified_model = replacer.replace_all(
>>>     filter_leaf,
>>>     "dilatedConvInt8",
>>>     False,
>>>     ccl_func_template_params=["dilations", "kernel_shape", "fixed_point_frac_bits"],
>>>     fixed_point_frac_bits=16,
>>>     match_callback=process_attributes
>>> )

fold_quant_nodes_into_custom_ops()

Fold quantization nodes into custom operators.

This method identifies quantization-related nodes (QuantizeLinear, DequantizeLinear) connected to custom operators and folds them into the custom operators. This improves performance by avoiding unnecessary quantization/dequantization operations.

This optimization is particularly useful for quantized models where quantization nodes are often inserted at the inputs and outputs of computational blocks.

Returns:

• None The method modifies the model in-place.

Notes:

The method performs two main optimizations:

1. DequantizeLinear folding at inputs:

   • Identifies DequantizeLinear nodes that feed into custom operators
   • Moves the dequantization into the custom operator's subgraph
   • Updates the custom operator to take quantized input directly

2. QuantizeLinear folding at outputs:

   • Identifies QuantizeLinear nodes that consume outputs from custom operators
   • Moves the quantization into the custom operator's subgraph
   • Updates the custom operator to produce quantized output directly

This optimization is only applied when it's safe - if a quantization node has multiple consumers, it won't be folded to maintain correctness.

Examples:

>>> replacer = CustomOpReplacer(model)
>>> replacer.fold_quant_nodes_into_custom_ops()