diff --git a/onnxscript/optimizer/__init__.py b/onnxscript/optimizer/__init__.py index d64a69d1c..e643762aa 100644 --- a/onnxscript/optimizer/__init__.py +++ b/onnxscript/optimizer/__init__.py @@ -129,6 +129,20 @@ def optimize_ir( input_size_limit: int = _DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT, output_size_limit: int = _DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT, ) -> None: + """ + Optimizes a model. + + Args: + 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 + 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: Not supported currently (has no effect). Meant to stop the + outer optimization loop if no change is detected in one iteration. + """ del stop_if_no_change # Looks like rewriter doesn't support this yet. _inliner.inline(model) for _ in range(num_iterations): diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 3e210ce59..5d76379cb 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -707,7 +707,9 @@ def process_node(self, node: ir.Node): if logger.isEnabledFor(logging.DEBUG): input_sizes = [input.size for input in input_values] logger.debug( - f"Skipping constant folding for op {node.op_type} due to large input size: {input_sizes}" + "Skipping constant folding for op %s due to large input size: %s", + node.op_type, + input_sizes, ) return None