lint

onnxscript/function_libs/torch_lib/ops/common.py

@@ -5,7 +5,8 @@
 # mypy: disable-error-code="misc,arg-type,type-arg,valid-type,assignment,return-value"
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import numpy.typing as npt
+import onnx
 
 import onnxscript
 import onnxscript.values
@@ -15,9 +16,6 @@
 from onnxscript.function_libs.torch_lib.tensor_typing import RealType
 from onnxscript.onnx_types import COMPLEX64, COMPLEX128, DOUBLE, FLOAT, TensorType
 
-if TYPE_CHECKING:
-    import onnx
-
 COMPLEX64_TYPE = COMPLEX64.dtype
 COMPLEX128_TYPE = COMPLEX128.dtype
 
@@ -66,6 +64,17 @@ def cast_to(a: RealType, dtype: int) -> RealType:
     return result
 
 
-def constant(array, dtype: int | onnx.TensorProto.DataType | ir.DataType) -> TensorType:
-    """Utility for creating a constant tensor."""
+def constant(
+    array: npt.ArrayLike | onnx.TensorProto | ir.DLPackCompatible | ir.ArrayCompatible,
+    dtype: int | onnx.TensorProto.DataType | ir.DataType,
+) -> TensorType:
+    """Utility for creating a constant tensor.
+
+    Args:
+        array: The array to convert to a constant tensor.
+        dtype: The data type of the tensor.
+
+    Returns:
+        A constant node.
+    """
     return op.Constant(value=ir.tensor(array, dtype=ir.DataType(dtype)))

tests/function_libs/torch_lib/quantization_test.py

@@ -12,8 +12,14 @@
 from torch.ao.quantization import quantize_pt2e
 from torch.ao.quantization.quantizer import xnnpack_quantizer
 
+from onnxscript._internal import version_utils
+
 
 class QuantizedModelExportTest(unittest.TestCase):
+    @unittest.skipIf(
+        version_utils.torch_older_than("2.4"),
+        "Dynamo exporter fails at the modularization step.",
+    )
     def test_simple_quantized_model(self):
         class TestModel(torch.nn.Module):
             def __init__(self):