-
Notifications
You must be signed in to change notification settings - Fork 53
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[API] Create stable APIs for PyTorch 2.6
optimize is turned on. It will be controlled by an option in PyTorch
- Loading branch information
1 parent
db30dbb
commit c73ae34
Showing
2 changed files
with
52 additions
and
46 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,26 @@ | ||
# Copyright (c) Microsoft Corporation. | ||
# Licensed under the MIT License. | ||
"""Stable APIs for PyTorch 2.6.""" | ||
|
||
from __future__ import annotations | ||
|
||
__all__ = [ | ||
"check_model", | ||
"convert_version", | ||
"get_torchlib_ops", | ||
"optimize", | ||
"save_model_with_external_data", | ||
] | ||
from onnxscript import ir, optimizer | ||
from onnxscript._framework_apis.torch_2_5 import ( | ||
check_model, | ||
convert_version, | ||
get_torchlib_ops, | ||
save_model_with_external_data, | ||
) | ||
|
||
|
||
def optimize(model: ir.Model) -> ir.Model: | ||
"""Optimize the model.""" | ||
optimizer.optimize_ir(model) | ||
return model |