[ONNX] Remove beartype usage (pytorch#130484)

beartype has served us well in identifying type errors and ensuring we call internal functions with the correct arguments (thanks!). However, the value of having beartype is diminished because of the following:

1. When beartype improves support for better Dict[] type checking, it discovered typing mistakes in some functions that were previously uncaught. This caused the exporter to fail with newer versions beartype when it used to succeed. Since we cannot fix PyTorch and release a new version just because of this, it creates confusion for users that have beartype in their environment from using torch.onnx
2. beartype adds an additional call line in the traceback, which makes the already thick dynamo stack even larger, affecting readability when users diagnose errors with the traceback.
3. Since the typing annotations need to be evaluated, we cannot use new syntaxes like `|` because we need to maintain compatibility with Python 3.8. We don't want to wait for PyTorch take py310 as the lowest supported Python before using the new typing syntaxes.

Pull Request resolved: pytorch#130484
Approved by: https://github.com/titaiwangms

Author: justinchuby

.ci/docker/common/install_onnx.sh

@@ -10,7 +10,6 @@ retry () {
 
 # A bunch of custom pip dependencies for ONNX
 pip_install \
-  beartype==0.15.0 \
   filelock==3.9.0 \
   flatbuffers==2.0 \
   mock==5.0.1 \

test/onnx/dynamo/test_exporter_api.py

@@ -3,18 +3,16 @@
 import os
 
 import onnx
-from beartype import roar
 
 import torch
 from torch.onnx import dynamo_export, ExportOptions, ONNXProgram
-from torch.onnx._internal import exporter, io_adapter
+from torch.onnx._internal import exporter
 from torch.onnx._internal.exporter import (
     LargeProtobufONNXProgramSerializer,
     ONNXProgramSerializer,
     ProtobufONNXProgramSerializer,
     ResolvedExportOptions,
 )
-from torch.onnx._internal.fx import diagnostics
 
 from torch.testing._internal import common_utils
 
@@ -49,15 +47,6 @@ def forward(self, x):
 
 
 class TestExportOptionsAPI(common_utils.TestCase):
-    def test_raise_on_invalid_argument_type(self):
-        expected_exception_type = roar.BeartypeException
-        with self.assertRaises(expected_exception_type):
-            ExportOptions(dynamic_shapes=2)  # type: ignore[arg-type]
-        with self.assertRaises(expected_exception_type):
-            ExportOptions(diagnostic_options="DEBUG")  # type: ignore[arg-type]
-        with self.assertRaises(expected_exception_type):
-            ResolvedExportOptions(options=12)  # type: ignore[arg-type]
-
     def test_dynamic_shapes_default(self):
         options = ResolvedExportOptions(ExportOptions())
         self.assertFalse(options.dynamic_shapes)
@@ -120,7 +109,6 @@ def test_save_to_file_using_specified_serializer_without_inheritance(self):
 
         # NOTE: Inheritance from `ONNXProgramSerializer` is not required.
         # Because `ONNXProgramSerializer` is a Protocol class.
-        # `beartype` will not complain.
         class CustomSerializer:
             def serialize(
                 self, onnx_program: ONNXProgram, destination: io.BufferedIOBase
@@ -196,27 +184,8 @@ def test_raise_from_diagnostic_warning_when_diagnostic_option_warning_as_error_i
                 ),
             )
 
-    def test_raise_on_invalid_save_argument_type(self):
-        with self.assertRaises(roar.BeartypeException):
-            ONNXProgram(torch.nn.Linear(2, 3))  # type: ignore[arg-type]
-        onnx_program = ONNXProgram(
-            onnx.ModelProto(),
-            io_adapter.InputAdapter(),
-            io_adapter.OutputAdapter(),
-            diagnostics.DiagnosticContext("test", "1.0"),
-            fake_context=None,
-        )
-        with self.assertRaises(roar.BeartypeException):
-            onnx_program.save(None)  # type: ignore[arg-type]
-        onnx_program.model_proto
-
 
 class TestProtobufONNXProgramSerializerAPI(common_utils.TestCase):
-    def test_raise_on_invalid_argument_type(self):
-        with self.assertRaises(roar.BeartypeException):
-            serializer = ProtobufONNXProgramSerializer()
-            serializer.serialize(None, None)  # type: ignore[arg-type]
-
     def test_serialize_raises_when_model_greater_than_2gb(self):
         onnx_program = torch.onnx.dynamo_export(_LargeModel(), torch.randn(1))
         serializer = ProtobufONNXProgramSerializer()

test/onnx/internal/test_beartype.py

@@ -34,7 +34,6 @@
 import torch
 from torch import export as torch_export
 from torch.onnx import _constants, verification
-from torch.onnx._internal import _beartype
 from torch.onnx._internal.fx import diagnostics
 from torch.testing._internal import common_utils
 from torch.testing._internal.opinfo import core as opinfo_core
@@ -206,7 +205,6 @@ def _run_test(m, remained_onnx_input_idx, flatten=True, ignore_none=True):
         if not is_model_script and not self.is_script:
             _run_test(model, tracing_remained_onnx_input_idx)
 
-    @_beartype.beartype
     def run_test_with_fx_to_onnx_exporter_and_onnx_runtime(
         self,
         model: _ModelType,
@@ -360,7 +358,6 @@ def run_test_with_fx_to_onnx_exporter_and_onnx_runtime(
                 )
 
 
-@_beartype.beartype
 def run_ort(
     onnx_model: Union[str, torch.onnx.ONNXProgram],
     pytorch_inputs: Sequence[_InputArgsType],
@@ -406,7 +403,6 @@ def run_ort(
     return session.run(None, ort_input)
 
 
-@_beartype.beartype
 def _try_clone_model(model: _ModelType) -> _ModelType:
     """Used for preserving original model in case forward mutates model states."""
     try:
@@ -418,14 +414,12 @@ def _try_clone_model(model: _ModelType) -> _ModelType:
         return model
 
 
-@_beartype.beartype
 def _try_clone_inputs(input_args, input_kwargs):
     ref_input_args = copy.deepcopy(input_args)
     ref_input_kwargs = copy.deepcopy(input_kwargs)
     return ref_input_args, ref_input_kwargs
 
 
-@_beartype.beartype
 def _compare_pytorch_onnx_with_ort(
     onnx_program: torch.onnx.ONNXProgram,
     model: _ModelType,

test/onnx/onnx_test_common.py

@@ -22,7 +22,7 @@
 from torch import nn
 
 from torch._subclasses import fake_tensor
-from torch.onnx._internal import _beartype, exporter
+from torch.onnx._internal import exporter
 from torch.onnx._internal.fx import (
     diagnostics,
     fx_symbolic_graph_extractor,
@@ -721,7 +721,6 @@ def forward(self, x):
             CustomModule(), (torch.randn(1, 2, 3),)
         )
 
-    @_beartype.beartype
     def _test_fx_symbolic_tracer_large_scale_exporter(
         self,
         model_name: str,
@@ -955,7 +954,6 @@ def setUp(self):
         super().setUp()
         self.ort_version = onnxruntime.__version__
 
-    @_beartype.beartype
     def _test_fake_tensor_mode_exporter(
         self,
         model_name: str,

test/onnx/test_fx_to_onnx_with_onnxruntime.py

@@ -35,10 +35,6 @@ Checkout topk in symbolic_opset10.py, and upsample_nearest2d in symbolic_opset8.
   `_jit_pass_onnx_remove_inplace_ops_for_onnx`, and
   transparently dispatched to their non inplace versions in
   "run_symbolic_function". See Note [Export inplace](#export-inplace)
-- Required: Annotate new symbolic functions with type annotations and decorate
-  with `@_beartype.beartype` to enable runtime type checking.
-  `@_beartype.beartype` should typically be the closest to the function to
-  ensure proper type checking.
 
 ### A note on Tensor types