Expose the should_fold option to optimize() (#2594)

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

onnxscript/optimizer/_optimizer.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+from typing import Callable
 
 import onnx_ir as ir
 import onnx_ir.passes.common as common_passes
@@ -21,6 +22,7 @@ def optimize_ir(
     stop_if_no_change: bool = True,
     input_size_limit: int = _constant_folding.DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT,
     output_size_limit: int = _constant_folding.DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT,
+    should_fold: Callable[[ir.Node], bool | None] = lambda node: None,
     inline: bool = True,
 ) -> None:
     """Optimizes a model.
@@ -29,11 +31,15 @@ def optimize_ir(
         model: The model to be optimized.
         num_iterations: Number of times the optimization loop is repeated.
         onnx_shape_inference: Applies node-level shape-inference as part of optimization
+        stop_if_no_change: Stop the optimization loop if no change is detected in an iteration.
         input_size_limit: Will not apply constant folding to ops with any input of size
             greater than this. Does not apply to special ops like Shape() and Size().
         output_size_limit: Will not rewrite any foldable-op into a Constant op if the size
             of the output tensor is greater than this.
-        stop_if_no_change: Stop the optimization loop if no change is detected in an iteration.
+        should_fold: An optional function that takes a node and returns True if
+            the node should be considered for folding.
+            The function should return True/False value to indicate if this particular
+            node should be folded, or None to use the default folding rules.
         inline: If True, inlines all functions in the model.
     """
     passes = [
@@ -43,6 +49,7 @@ def optimize_ir(
                     shape_inference=onnx_shape_inference,
                     input_size_limit=input_size_limit,
                     output_size_limit=output_size_limit,
+                    should_fold=should_fold,
                 ),
                 rewriter.RewritePass(rewriter._DEFAULT_REWRITE_RULES),
                 common_passes.RemoveUnusedNodesPass(),