Fix misleading annotation in the documentation (#2046)

The function does not seem to work inplace in all cases.

---------

Co-authored-by: Ti-Tai Wang <[email protected]>
Co-authored-by: Justin Chu <[email protected]>

Author: 3 people

docs/conf.py

@@ -88,7 +88,10 @@
     "python": (f"https://docs.python.org/{sys.version_info.major}", None),
     "matplotlib": ("https://matplotlib.org/stable/", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
+    "onnx": ("https://onnx.ai/onnx/", None),
     "onnxruntime": ("https://onnxruntime.ai/docs/api/python/", None),
+    "scipy": ("https://docs.scipy.org/doc/scipy/", None),
+    "torch": ("https://pytorch.org/docs/main/", None),
 }
 
 # -- Options for Sphinx Gallery ----------------------------------------------

onnxscript/optimizer/__init__.py

@@ -2,6 +2,15 @@
 # Licensed under the MIT License.
 from __future__ import annotations
 
+__all__ = [
+    "fold_constants",
+    "fold_constants_ir",
+    "remove_unused_nodes",
+    "optimize",
+    "optimize_ir",
+    "basic_constant_propagation",
+]
+
 import onnx
 
 import onnxscript.optimizer._constant_folding as constant_folding
@@ -15,25 +24,17 @@
 fold_constants_ir = constant_folding.fold_constants
 
 
-def optimize(model: ir.Model | onnx.ModelProto, *args, **kwargs):
+def optimize(model: ir.Model, *args, **kwargs) -> ir.Model:
     if isinstance(model, ir.Model):
-        return optimize_ir(model, *args, **kwargs)
+        # In that case, this is done inplace.
+        optimize_ir(model, *args, **kwargs)
+        return model
     else:
         return legacy_optimizer.optimize(model, *args, **kwargs)
 
 
-def fold_constants(model: ir.Model | onnx.ModelProto, *args, **kwargs):
+def fold_constants(model: ir.Model | onnx.ModelProto, *args, **kwargs) -> bool:
     if isinstance(model, ir.Model):
         return constant_folding.fold_constants(model, *args, **kwargs)
     else:
         return legacy_constant_folding.fold_constants(model, *args, **kwargs)
-
-
-__all__ = [
-    "fold_constants",
-    "fold_constants_ir",
-    "remove_unused_nodes",
-    "optimize",
-    "optimize_ir",
-    "basic_constant_propagation",
-]

onnxscript/tools/benchmark/benchmark_helpers.py

@@ -450,11 +450,12 @@ def optimize_model_proto(
         begin = time.perf_counter()
 
         if value == "optimize":
-            model_proto = onnxscript.optimizer.optimize(
-                model_proto,
+            model_ir = onnxscript.optimizer.optimize(
+                ir.from_proto(model_proto),
                 num_iterations=2,
                 onnx_shape_inference=False,
             )
+            model_proto = ir.to_proto(model_ir)
 
         elif value == "rewrite":
             model_proto = onnxscript.rewriter.rewrite(model_proto)

onnxscript/tools/transformers_models/__init__.py

@@ -42,14 +42,14 @@ def export_to_onnx(
         else:
             prog = torch.onnx.dynamo_export(model, *args)
     assert prog is not None
-    model_proto = prog.model_proto
+    model = prog.model
     if optimize:
-        model_proto = onnxscript.optimizer.optimize(
-            model_proto,
+        model = onnxscript.optimizer.optimize(
+            model,
             num_iterations=2,
-            onnx_shape_inference=True,
         )
-        model_proto = onnxscript.rewriter.rewrite(model_proto)
+        model = onnxscript.rewriter.rewrite(model)
+        model_proto = onnxscript.ir.to_proto(model)
         model_proto = onnx.inliner.inline_local_functions(model_proto)
     return model_proto