update typing

Author: justinchuby

onnxscript/ir/passes/common/_c_api_utils.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Callable, TypeVar
 
 from onnxscript import ir
 
@@ -17,19 +17,18 @@
 # Temporarily remove initializers larger than this size to keep model size down
 # for the onnx.shape_inference call because it needs to serialize the model
 _BIG_TENSOR_SIZE_LIMIT = 1000  # 1KB
+_R = TypeVar("_R")
 
 
-def call_onnx_api(
-    func: Callable[[onnx.ModelProto], onnx.ModelProto], model: ir.Model
-) -> onnx.ModelProto:
+def call_onnx_api(func: Callable[[onnx.ModelProto], _R], model: ir.Model) -> _R:
     """Call an ONNX C API function by temporarily removing initializers.
 
     This is necessary because the ONNX C API does not support large models
     with initializers that have large tensor values. The input model is left
     unchanged no matter the call succeeds or not.
 
     Args:
-        func: Partially applied function that takes a model proto and returns a model proto.
+        func: Partially applied function that takes a model proto and returns anything.
         model: The IR model to pass to the API function.
 
     Returns:
@@ -61,7 +60,7 @@ def call_onnx_api(
 
     try:
         proto = ir.serde.serialize_model(model)
-        result_proto = func(proto)
+        result = func(proto)
     finally:
         # Restore the original initializer values so the model is unchanged
         for initializer in initializer_values:
@@ -73,4 +72,4 @@ def call_onnx_api(
         model.graph.inputs.clear()
         model.graph.inputs.extend(inputs)
 
-    return result_proto
+    return result

onnxscript/ir/passes/common/onnx_checker.py

@@ -39,15 +39,14 @@ def __init__(
     def call(self, model: ir.Model) -> ir.passes.PassResult:
         """Run the onnx checker on the model."""
 
-        def _partial_check_model(proto: onnx.ModelProto) -> onnx.ModelProto:
+        def _partial_check_model(proto: onnx.ModelProto) -> None:
             """Partial function to check the model."""
             onnx.checker.check_model(
                 proto,
                 full_check=self.full_check,
                 skip_opset_compatibility_check=self.skip_opset_compatibility_check,
                 check_custom_domain=self.check_custom_domain,
             )
-            return proto
 
         _c_api_utils.call_onnx_api(func=_partial_check_model, model=model)
         # The model is not modified