Expose the register as public register_constant_folder API (fix for Issue #2507)

onnxscript/optimizer/__init__.py

@@ -12,6 +12,7 @@
     "inline",
     "optimize_ir",
     "optimize",
+    "register_constant_folder",
     "remove_unused_nodes",
 ]
 
@@ -20,8 +21,14 @@
 
 import onnxscript.optimizer._constant_folding as constant_folding
 from onnxscript import ir
-from onnxscript.optimizer._constant_folding import FOLDED_FROM_KEY, basic_constant_propagation
+from onnxscript.optimizer._constant_folding import (
+    FOLDED_FROM_KEY,
+    basic_constant_propagation,
+)
 from onnxscript.optimizer._constant_folding import fold_constants as fold_constants_ir
+from onnxscript.optimizer._constant_folding import (
+    register as register_constant_folder,
+)
 from onnxscript.optimizer._optimizer import optimize_ir
 
 _ModelProtoOrIr = TypeVar("_ModelProtoOrIr", onnx.ModelProto, ir.Model)

onnxscript/optimizer/_constant_folding.py

@@ -258,6 +258,32 @@ def lookup_evaluators(self, domain: str, opname: str, version: int):
     def register(
         self, opname: str, domain: str = "", version=None
     ) -> Callable[[PartialEvaluatorFunction], PartialEvaluatorFunction]:
+        """Register a custom constant folding rule for an operator.
+
+        The decorated function must have the signature:
+        `(node: ir.Node, op: OptimizerContext, state: OptimizerState) -> ReturnValue`
+
+        Args:
+            opname: The name of the operator to fold (e.g., "Add").
+            domain: The domain of the operator. Defaults to "" (standard ONNX).
+            version: The opset version or version range for which this rule is valid.
+                Can be an integer (e.g., 13) or a tuple of (min_version, max_version)
+                where None indicates no limit (e.g., (13, None)).
+
+        Returns:
+            A decorator that registers the function as a constant folder for the op.
+
+        Example:
+
+            .. code-block:: python
+
+                from onnxscript.optimizer import register_constant_folder
+
+                @register_constant_folder("CustomOp", domain="my.domain")
+                def fold_custom_op(node: ir.Node, op: OptimizerContext, state: OptimizerState) -> ReturnValue:
+                    # custom folding logic
+                    return op.Constant(value_int=42)
+        """
         if (domain, opname) in self.op_evaluators:
             evaluator_list = self.op_evaluators[(domain, opname)]
         else:

onnxscript/optimizer/_constant_folding_test.py

@@ -879,6 +879,37 @@ def test_initializer_as_graph_output_is_not_removed(self):
         self.assertIn("y", output_names)
         self.assertIn("z", output_names)
 
+    def test_register_constant_folder(self):
+        @optimizer.register_constant_folder("CustomAdd", domain="test.custom")
+        def fold_custom_add(node: ir.Node, op, state: _constant_folding.OptimizerState):
+            return op.Constant(value_int=42)
+
+        try:
+            model = ir.from_onnx_text(
+                """
+                <ir_version: 7, opset_import: [ "" : 17, "test.custom" : 1 ]>
+                agraph (float[N] x) => (int64 z) {
+                    z = test.custom.CustomAdd(x)
+                }
+                """
+            )
+
+            optimized = self._fold(model)
+            custom_add_nodes = [n for n in optimized.graph if n.op_type == "CustomAdd"]
+            self.assertEqual(len(custom_add_nodes), 0)
+            z_value = optimized.graph.outputs[0]
+            self.assertIsNotNone(z_value.const_value)
+            np.testing.assert_equal(z_value.const_value.numpy(), np.array(42, dtype=np.int64))
+        finally:
+            key = ("test.custom", "CustomAdd")
+            evaluators = _constant_folding.registry.op_evaluators.get(key)
+            if evaluators is not None:
+                _constant_folding.registry.op_evaluators[key] = [
+                    e for e in evaluators if e.function is not fold_custom_add
+                ]
+                if not _constant_folding.registry.op_evaluators[key]:
+                    del _constant_folding.registry.op_evaluators[key]
+
 
 def _all_value_names_unique(model: ir.Model) -> bool:
     """Return True if all named values in the top-level graph have unique names."""