Revert "[ONNX] Remove beartype usage (pytorch#130484)"

This reverts commit 1794c35.

Reverted pytorch#130484 on behalf of https://github.com/clee2000 due to test_sympy_utils failure is real https://github.com/pytorch/pytorch/actions/runs/9961499559/job/27523758780 https://hud.pytorch.org/pytorch/pytorch/commit/1794c35912025aa44b0d70f67ff664b4f7bd1014.  Dr CI is matching with commits in current commit? ([comment](pytorch#130484 (comment)))

Author: pytorchmergebot

.ci/docker/common/install_onnx.sh

@@ -10,6 +10,7 @@ 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,16 +3,18 @@
 import os
 
 import onnx
+from beartype import roar
 
 import torch
 from torch.onnx import dynamo_export, ExportOptions, ONNXProgram
-from torch.onnx._internal import exporter
+from torch.onnx._internal import exporter, io_adapter
 from torch.onnx._internal.exporter import (
     LargeProtobufONNXProgramSerializer,
     ONNXProgramSerializer,
     ProtobufONNXProgramSerializer,
     ResolvedExportOptions,
 )
+from torch.onnx._internal.fx import diagnostics
 
 from torch.testing._internal import common_utils
 
@@ -47,6 +49,15 @@ 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)
@@ -109,6 +120,7 @@ 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
@@ -184,8 +196,27 @@ 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

@@ -0,0 +1,84 @@
+# Owner(s): ["module: onnx"]
+"""Unit tests for the internal beartype wrapper module."""
+
+import unittest
+
+from torch.onnx._internal import _beartype
+from torch.testing._internal import common_utils
+
+
+def beartype_installed():
+    try:
+        import beartype  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
+
+def skip_if_beartype_not_installed(test_case):
+    return unittest.skipIf(not beartype_installed(), "beartype is not installed")(
+        test_case
+    )
+
+
+def func_with_type_hint(x: int) -> int:
+    return x
+
+
+def func_with_incorrect_type_hint(x: int) -> str:
+    return x  # type: ignore[return-value]
+
+
+@common_utils.instantiate_parametrized_tests
+class TestBeartype(common_utils.TestCase):
+    def test_create_beartype_decorator_returns_no_op_decorator_when_disabled(self):
+        decorator = _beartype._create_beartype_decorator(
+            _beartype.RuntimeTypeCheckState.DISABLED,
+        )
+        decorated = decorator(func_with_incorrect_type_hint)
+        decorated("string_input")  # type: ignore[arg-type]
+
+    @skip_if_beartype_not_installed
+    def test_create_beartype_decorator_warns_when_warnings(self):
+        decorator = _beartype._create_beartype_decorator(
+            _beartype.RuntimeTypeCheckState.WARNINGS,
+        )
+        decorated = decorator(func_with_incorrect_type_hint)
+        with self.assertWarns(_beartype.CallHintViolationWarning):
+            decorated("string_input")  # type: ignore[arg-type]
+
+    @common_utils.parametrize("arg", [1, "string_input"])
+    @skip_if_beartype_not_installed
+    def test_create_beartype_decorator_errors_when_errors(self, arg):
+        import beartype
+
+        decorator = _beartype._create_beartype_decorator(
+            _beartype.RuntimeTypeCheckState.ERRORS,
+        )
+        decorated = decorator(func_with_incorrect_type_hint)
+        with self.assertRaises(beartype.roar.BeartypeCallHintViolation):
+            decorated(arg)
+
+    @skip_if_beartype_not_installed
+    def test_create_beartype_decorator_warning_calls_function_once(self):
+        call_count = 0
+
+        def func_with_incorrect_type_hint_and_side_effect(x: int) -> str:
+            nonlocal call_count
+            call_count += 1
+            return x  # type: ignore[return-value]
+
+        decorator = _beartype._create_beartype_decorator(
+            _beartype.RuntimeTypeCheckState.WARNINGS,
+        )
+        decorated = decorator(func_with_incorrect_type_hint_and_side_effect)
+        decorated("string_input")  # type: ignore[arg-type]
+        self.assertEqual(call_count, 1)
+        decorated(1)
+        # The return value violates the type hint, but the function is called
+        # only once.
+        self.assertEqual(call_count, 2)
+
+
+if __name__ == "__main__":
+    common_utils.run_tests()

test/onnx/onnx_test_common.py

@@ -34,6 +34,7 @@
 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
@@ -205,6 +206,7 @@ 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,
@@ -358,6 +360,7 @@ 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],
@@ -403,6 +406,7 @@ 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:
@@ -414,12 +418,14 @@ 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/test_fx_to_onnx_with_onnxruntime.py

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

torch/onnx/README.md

@@ -35,6 +35,10 @@ 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
 

torch/onnx/_internal/_beartype.py

@@ -0,0 +1,132 @@
+# mypy: allow-untyped-defs
+"""An internal wrapper for the beartype library.
+
+The module returns a no-op decorator when the beartype library is not installed.
+"""
+import enum
+import functools
+import os
+import traceback
+import typing
+import warnings
+from types import ModuleType
+
+try:
+    import beartype as _beartype_lib  # type: ignore[import]
+    from beartype import roar as _roar  # type: ignore[import]
+
+    # Beartype warns when we import from typing because the types are deprecated
+    # in Python 3.9. But there will be a long time until we can move to using
+    # the native container types for type annotations (when 3.9 is the lowest
+    # supported version). So we silence the warning.
+    warnings.filterwarnings(
+        "ignore",
+        category=_roar.BeartypeDecorHintPep585DeprecationWarning,
+    )
+
+    if _beartype_lib.__version__ == "0.16.0":
+        # beartype 0.16.0 has a bug that causes it to crash when used with
+        # PyTorch. See https://github.com/beartype/beartype/issues/282
+        warnings.warn("beartype 0.16.0 is not supported. Please upgrade to 0.16.1+.")
+        _beartype_lib = None  # type: ignore[assignment]
+except ImportError:
+    _beartype_lib = None  # type: ignore[assignment]
+except Exception as e:
+    # Warn errors that are not import errors (unexpected).
+    warnings.warn(f"{e}")
+    _beartype_lib = None  # type: ignore[assignment]
+
+
+@enum.unique
+class RuntimeTypeCheckState(enum.Enum):
+    """Runtime type check state."""
+
+    # Runtime type checking is disabled.
+    DISABLED = enum.auto()
+    # Runtime type checking is enabled but warnings are shown only.
+    WARNINGS = enum.auto()
+    # Runtime type checking is enabled.
+    ERRORS = enum.auto()
+
+
+class CallHintViolationWarning(UserWarning):
+    """Warning raised when a type hint is violated during a function call."""
+
+    pass
+
+
+def _no_op_decorator(func):
+    return func
+
+
+def _create_beartype_decorator(
+    runtime_check_state: RuntimeTypeCheckState,
+):
+    # beartype needs to be imported outside of the function and aliased because
+    # this module overwrites the name "beartype".
+
+    if runtime_check_state == RuntimeTypeCheckState.DISABLED:
+        return _no_op_decorator
+    if _beartype_lib is None:
+        # If the beartype library is not installed, return a no-op decorator
+        return _no_op_decorator
+
+    assert isinstance(_beartype_lib, ModuleType)
+
+    if runtime_check_state == RuntimeTypeCheckState.ERRORS:
+        # Enable runtime type checking which errors on any type hint violation.
+        return _beartype_lib.beartype
+
+    # Warnings only
+    def beartype(func):
+        """Warn on type hint violation."""
+
+        if "return" in func.__annotations__:
+            # Remove the return type from the func function's
+            # annotations so that the beartype decorator does not complain
+            # about the return type.
+            return_type = func.__annotations__["return"]
+            del func.__annotations__["return"]
+            beartyped = _beartype_lib.beartype(func)
+            # Restore the return type to the func function's annotations
+            func.__annotations__["return"] = return_type
+        else:
+            beartyped = _beartype_lib.beartype(func)
+
+        @functools.wraps(func)
+        def _coerce_beartype_exceptions_to_warnings(*args, **kwargs):
+            try:
+                return beartyped(*args, **kwargs)
+            except _roar.BeartypeCallHintParamViolation:
+                # Fall back to the original function if the beartype hint is violated.
+                warnings.warn(
+                    traceback.format_exc(),
+                    category=CallHintViolationWarning,
+                    stacklevel=2,
+                )
+
+            return func(*args, **kwargs)  # noqa: B012
+
+        return _coerce_beartype_exceptions_to_warnings
+
+    return beartype
+
+
+if typing.TYPE_CHECKING:
+    # This is a hack to make mypy play nicely with the beartype decorator.
+    def beartype(func):
+        return func
+
+else:
+    _TORCH_ONNX_EXPERIMENTAL_RUNTIME_TYPE_CHECK = os.getenv(
+        "TORCH_ONNX_EXPERIMENTAL_RUNTIME_TYPE_CHECK"
+    )
+    if _TORCH_ONNX_EXPERIMENTAL_RUNTIME_TYPE_CHECK == "ERRORS":
+        _runtime_type_check_state = RuntimeTypeCheckState.ERRORS
+    elif _TORCH_ONNX_EXPERIMENTAL_RUNTIME_TYPE_CHECK == "DISABLED":
+        _runtime_type_check_state = RuntimeTypeCheckState.DISABLED
+    else:
+        _runtime_type_check_state = RuntimeTypeCheckState.WARNINGS
+    beartype = _create_beartype_decorator(_runtime_type_check_state)
+    # Make sure that the beartype decorator is enabled whichever path we took.
+    assert beartype is not None