Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion onnxscript/rewriter/ort_fusions/_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,7 @@ def optimize_for_ort(
config_name: str | None = None,
*,
debug: bool = False,
clear_metadata: bool = False,
) -> tuple[ir.Model, dict[str, int]]:
"""
Optimize the model for ORT backend.
Expand All @@ -128,6 +129,7 @@ def optimize_for_ort(
Typically it identifies the Execution Provider (EP) to optimize for.
If None, the default configuration will be used.
debug: If debug is True, enable pattern matching tracer for debugging.
clear_metadata: If True, clear metadata and doc strings from the model.

Returns:
A tuple containing:
Expand All @@ -145,7 +147,6 @@ def optimize_for_ort(
passes = ir.passes.Sequential(
# Apply the ORT optimization passes.
# https://github.com/microsoft/onnxruntime/blob/74dcf7e296639095dfa55d31336998b6f719ed76/onnxruntime/python/tools/transformers/dynamo_onnx_helper.py#L172
common_passes.ClearMetadataAndDocStringPass(),
# https://github.com/microsoft/onnxruntime/blob/74dcf7e296639095dfa55d31336998b6f719ed76/onnxruntime/python/tools/transformers/dynamo_onnx_helper.py#L139
common_passes.LiftConstantsToInitializersPass(lift_all_constants=False, size_limit=1),
common_passes.RemoveInitializersFromInputsPass(),
Expand All @@ -154,4 +155,8 @@ def optimize_for_ort(
assert passes.in_place
result = passes(model)
assert result.model is model

if clear_metadata:
common_passes.ClearMetadataAndDocStringPass()(model)

return model, fusion_count
Loading