From 1cb6a5858251ad956ef493f4e08fcf21207866e2 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 8 Jan 2023 04:09:21 +0000 Subject: [PATCH 1/9] fix(atenlib): Cast SymInt to INT64 to fix Windows CI --- onnxscript/function_libs/torch_aten/ops/core.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/onnxscript/function_libs/torch_aten/ops/core.py b/onnxscript/function_libs/torch_aten/ops/core.py index e9e45f7f15..bd74a44c2b 100644 --- a/onnxscript/function_libs/torch_aten/ops/core.py +++ b/onnxscript/function_libs/torch_aten/ops/core.py @@ -1586,10 +1586,10 @@ def aten_exp2(self: TFloat) -> TFloat: @torch_op("aten::expand") -def aten_expand(self: TTensor, size: INT64) -> TTensor: +def aten_expand(self: TTensor, size: TInt) -> TTensor: # expand(Tensor(a) self, SymInt[] size, *, bool implicit=False) -> Tensor(a) - size = op.Cast(size, to=INT64.dtype) # to INT64 + size = op.Cast(size, to=INT64.dtype) return op.Expand(self, size) @@ -3462,10 +3462,11 @@ def aten_new_empty_strided(self: TensorType, size: INT64, stride: INT64) -> Tens @torch_op("aten::new_full") def aten_new_full( - self, size: INT64, fill_value, dtype: int = FLOAT.dtype + self, size: TInt, fill_value, dtype: int = FLOAT.dtype ): # pylint: disable=unused-argument # new_full(Tensor self, SymInt[] size, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + size = op.Cast(size, to=INT64.dtype) fill_value = op.Cast(fill_value, to=dtype) return op.Expand(fill_value, size) @@ -4032,13 +4033,14 @@ def aten_renorm(self: TensorType, p: float, dim: int, maxnorm: float) -> TensorT @torch_op("aten::repeat") -def aten_repeat(self: TTensor, repeats: INT64) -> TTensor: +def aten_repeat(self: TTensor, repeats: TInt) -> TTensor: # repeat(Tensor self, SymInt[] repeats) -> Tensor if op.Size(repeats) == 0: result = self else: # TODO(justinchuby): Make ones_like a function when onnxscript supports it + repeats = op.Cast(repeats, to=INT64.dtype) # shape = ones_like(repeats) := { one = op.Constant(value_int=1) repeats_shape = op.Shape(repeats) @@ -4058,10 +4060,11 @@ def aten_repeat_interleave( @torch_op("aten::reshape") -def aten_reshape(self: TTensor, shape: INT64) -> TTensor: +def aten_reshape(self: TTensor, shape: TInt) -> TTensor: # reshape(Tensor(a) self, SymInt[] shape) -> Tensor(a) - shape = op.Cast(shape, to=INT64.dtype) # Reshape only support INT64 as 'shape' + # Reshape only support INT64 as 'shape' + shape = op.Cast(shape, to=INT64.dtype) return op.Reshape(self, shape) @@ -4919,7 +4922,7 @@ def aten_vdot(self: TensorType, other: TensorType) -> TensorType: @torch_op("aten::view") -def aten_view(self: TTensor, size: INT64) -> TTensor: +def aten_view(self: TTensor, size: TInt) -> TTensor: # view(Tensor(a) self, SymInt[] size) -> Tensor(a) size = op.Cast(size, to=INT64.dtype) # Reshape only support INT64 as second input From 5605b25c415de400330b1057f341d722df7325b0 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 8 Jan 2023 04:19:47 +0000 Subject: [PATCH 2/9] Short test results --- .github/workflows/main.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/main.yaml b/.github/workflows/main.yaml index 80240880e0..4b36c94f5a 100644 --- a/.github/workflows/main.yaml +++ b/.github/workflows/main.yaml @@ -77,7 +77,7 @@ jobs: pip list | grep torch - name: pytest - run: pytest -v onnxscript --cov=onnxscript --cov-report=xml -n auto + run: pytest -v onnxscript --cov=onnxscript --cov-report=xml -n=auto --tb=short - name: Install package run: pip install . From ef8ec9bd616d193076992eddfaf8b11b1812f441 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 8 Jan 2023 04:21:56 +0000 Subject: [PATCH 3/9] color output --- .github/workflows/main.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yaml b/.github/workflows/main.yaml index 4b36c94f5a..0ed32bcd69 100644 --- a/.github/workflows/main.yaml +++ b/.github/workflows/main.yaml @@ -77,14 +77,14 @@ jobs: pip list | grep torch - name: pytest - run: pytest -v onnxscript --cov=onnxscript --cov-report=xml -n=auto --tb=short + run: pytest -v onnxscript --cov=onnxscript --cov-report=xml -n=auto --tb=short --color=yes - name: Install package run: pip install . - name: Test examples if: ${{ matrix.test_examples }} - run: pytest -v docs/test + run: pytest -v docs/test -n=auto --color=yes - name: Build package run: python -m build From c1b3795233c10f42534ddcc00ecc358be0ccdc27 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 8 Jan 2023 04:31:16 +0000 Subject: [PATCH 4/9] pytest config --- .github/workflows/main.yaml | 4 ++-- pyproject.toml | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yaml b/.github/workflows/main.yaml index 0ed32bcd69..816047d9a4 100644 --- a/.github/workflows/main.yaml +++ b/.github/workflows/main.yaml @@ -77,14 +77,14 @@ jobs: pip list | grep torch - name: pytest - run: pytest -v onnxscript --cov=onnxscript --cov-report=xml -n=auto --tb=short --color=yes + run: pytest -v onnxscript --cov=onnxscript --cov-report=xml -n=auto - name: Install package run: pip install . - name: Test examples if: ${{ matrix.test_examples }} - run: pytest -v docs/test -n=auto --color=yes + run: pytest -v docs/test -n=auto - name: Build package run: python -m build diff --git a/pyproject.toml b/pyproject.toml index 0e529617d3..83526d7359 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ onnx = ["py.typed"] [tool.pytest.ini_options] filterwarnings = ["ignore::UserWarning", "ignore::DeprecationWarning"] +addopts = "-ra --tb=short --color=yes" [tool.mypy] follow_imports = "silent" # TODO: Remove when we fix all the mypy errors From 4c642206df02fd8dc62abb59b48a3bcebe478320 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 8 Jan 2023 04:31:22 +0000 Subject: [PATCH 5/9] ones and zeros --- onnxscript/function_libs/torch_aten/ops/core.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/onnxscript/function_libs/torch_aten/ops/core.py b/onnxscript/function_libs/torch_aten/ops/core.py index bd74a44c2b..85bef339c4 100644 --- a/onnxscript/function_libs/torch_aten/ops/core.py +++ b/onnxscript/function_libs/torch_aten/ops/core.py @@ -1589,7 +1589,7 @@ def aten_exp2(self: TFloat) -> TFloat: def aten_expand(self: TTensor, size: TInt) -> TTensor: # expand(Tensor(a) self, SymInt[] size, *, bool implicit=False) -> Tensor(a) - size = op.Cast(size, to=INT64.dtype) + size = op.Cast(size, to=6) return op.Expand(self, size) @@ -3530,9 +3530,10 @@ def aten_nuclear_norm(self: TensorType, keepdim: bool = False) -> TensorType: @torch_op("aten::ones") -def aten_ones(size: INT64, dtype: int = -1): +def aten_ones(size: TInt, dtype: int = -1) -> TTensor: # ones(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + size = op.Cast(size, to=INT64.dtype) one = op.Constant(value_float=1) if dtype != -1: one = op.Cast(one, to=dtype) @@ -4991,9 +4992,10 @@ def aten_xor(self: TensorType, other: TensorType) -> TensorType: @torch_op("aten::zeros") -def aten_zeros(size: INT64, dtype: int = -1): +def aten_zeros(size: TInt, dtype: int = -1) -> TTensor: # zeros(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + size = op.Cast(size, to=INT64.dtype) zero = op.Constant(value_float=0) if dtype != -1: zero = op.Cast(zero, to=dtype) From fee03deb39e74e0603da72869fc32976382ae36b Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 8 Jan 2023 04:36:43 +0000 Subject: [PATCH 6/9] oops --- onnxscript/function_libs/torch_aten/ops/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxscript/function_libs/torch_aten/ops/core.py b/onnxscript/function_libs/torch_aten/ops/core.py index 85bef339c4..d6509ded0b 100644 --- a/onnxscript/function_libs/torch_aten/ops/core.py +++ b/onnxscript/function_libs/torch_aten/ops/core.py @@ -1589,7 +1589,7 @@ def aten_exp2(self: TFloat) -> TFloat: def aten_expand(self: TTensor, size: TInt) -> TTensor: # expand(Tensor(a) self, SymInt[] size, *, bool implicit=False) -> Tensor(a) - size = op.Cast(size, to=6) + size = op.Cast(size, to=INT64.dtype) return op.Expand(self, size) From 6f2c9ea897080a1a3a8b119bba2c23b6c84d8285 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 8 Jan 2023 04:46:10 +0000 Subject: [PATCH 7/9] Fix types --- .../function_libs/torch_aten/ops/core.py | 23 +++++++++---------- onnxscript/function_libs/torch_aten/typing.py | 4 ++-- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/onnxscript/function_libs/torch_aten/ops/core.py b/onnxscript/function_libs/torch_aten/ops/core.py index d6509ded0b..c3c58ff000 100644 --- a/onnxscript/function_libs/torch_aten/ops/core.py +++ b/onnxscript/function_libs/torch_aten/ops/core.py @@ -16,6 +16,7 @@ from onnxscript import BOOL, DOUBLE, FLOAT, INT64 from onnxscript.function_libs.torch_aten.registration import torch_op from onnxscript.function_libs.torch_aten.typing import ( + IntType, TFloat, TFloatOrBFloat16, TInt, @@ -1586,7 +1587,7 @@ def aten_exp2(self: TFloat) -> TFloat: @torch_op("aten::expand") -def aten_expand(self: TTensor, size: TInt) -> TTensor: +def aten_expand(self: TTensor, size: IntType) -> TTensor: # expand(Tensor(a) self, SymInt[] size, *, bool implicit=False) -> Tensor(a) size = op.Cast(size, to=INT64.dtype) @@ -2178,7 +2179,7 @@ def aten_index_reduce( # FIXME(#277): Script when attributes can come before inputs @torch_op("aten::index_select", trace_only=True) -def aten_index_select(self: TTensor, dim: int, index: TInt) -> TTensor: +def aten_index_select(self: TTensor, dim: int, index: IntType) -> TTensor: # index_select(Tensor self, int dim, Tensor index) -> Tensor if op.Size(op.Shape(self)) == 0: @@ -3462,7 +3463,7 @@ def aten_new_empty_strided(self: TensorType, size: INT64, stride: INT64) -> Tens @torch_op("aten::new_full") def aten_new_full( - self, size: TInt, fill_value, dtype: int = FLOAT.dtype + self, size: IntType, fill_value, dtype: int = FLOAT.dtype ): # pylint: disable=unused-argument # new_full(Tensor self, SymInt[] size, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor @@ -3530,13 +3531,12 @@ def aten_nuclear_norm(self: TensorType, keepdim: bool = False) -> TensorType: @torch_op("aten::ones") -def aten_ones(size: TInt, dtype: int = -1) -> TTensor: +def aten_ones(size: IntType, dtype: int = FLOAT.dtype): # ones(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor size = op.Cast(size, to=INT64.dtype) one = op.Constant(value_float=1) - if dtype != -1: - one = op.Cast(one, to=dtype) + one = op.Cast(one, to=dtype) return op.Expand(one, size) @@ -4034,7 +4034,7 @@ def aten_renorm(self: TensorType, p: float, dim: int, maxnorm: float) -> TensorT @torch_op("aten::repeat") -def aten_repeat(self: TTensor, repeats: TInt) -> TTensor: +def aten_repeat(self: TTensor, repeats: IntType) -> TTensor: # repeat(Tensor self, SymInt[] repeats) -> Tensor if op.Size(repeats) == 0: @@ -4061,7 +4061,7 @@ def aten_repeat_interleave( @torch_op("aten::reshape") -def aten_reshape(self: TTensor, shape: TInt) -> TTensor: +def aten_reshape(self: TTensor, shape: IntType) -> TTensor: # reshape(Tensor(a) self, SymInt[] shape) -> Tensor(a) # Reshape only support INT64 as 'shape' @@ -4923,7 +4923,7 @@ def aten_vdot(self: TensorType, other: TensorType) -> TensorType: @torch_op("aten::view") -def aten_view(self: TTensor, size: TInt) -> TTensor: +def aten_view(self: TTensor, size: IntType) -> TTensor: # view(Tensor(a) self, SymInt[] size) -> Tensor(a) size = op.Cast(size, to=INT64.dtype) # Reshape only support INT64 as second input @@ -4992,13 +4992,12 @@ def aten_xor(self: TensorType, other: TensorType) -> TensorType: @torch_op("aten::zeros") -def aten_zeros(size: TInt, dtype: int = -1) -> TTensor: +def aten_zeros(size: IntType, dtype: int = FLOAT.dtype): # zeros(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor size = op.Cast(size, to=INT64.dtype) zero = op.Constant(value_float=0) - if dtype != -1: - zero = op.Cast(zero, to=dtype) + zero = op.Cast(zero, to=dtype) return op.Expand(zero, size) diff --git a/onnxscript/function_libs/torch_aten/typing.py b/onnxscript/function_libs/torch_aten/typing.py index 8ab89c647a..9587a8534d 100644 --- a/onnxscript/function_libs/torch_aten/typing.py +++ b/onnxscript/function_libs/torch_aten/typing.py @@ -41,7 +41,7 @@ UINT8, ] _FloatType = Union[FLOAT16, FLOAT, DOUBLE] -_IntType = Union[INT8, INT16, INT32, INT64] +IntType = Union[INT8, INT16, INT32, INT64] _RealType = Union[ BFLOAT16, FLOAT16, @@ -56,7 +56,7 @@ TTensor = TypeVar("TTensor", bound=_TensorType) TFloat = TypeVar("TFloat", bound=_FloatType) TFloatOrBFloat16 = TypeVar("TFloatOrBFloat16", bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16]) -TInt = TypeVar("TInt", bound=_IntType) +TInt = TypeVar("TInt", bound=IntType) TReal = TypeVar("TReal", bound=_RealType) TRealUnlessInt16OrInt8 = TypeVar( "TRealUnlessInt16OrInt8", bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16, INT32, INT64] From 3801ec591f42a6f1fcabb135358840d62533680e Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 9 Jan 2023 16:20:05 +0000 Subject: [PATCH 8/9] TInt --- onnxscript/function_libs/torch_aten/ops/core.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxscript/function_libs/torch_aten/ops/core.py b/onnxscript/function_libs/torch_aten/ops/core.py index 64b1a8306a..9db8b7ab2f 100644 --- a/onnxscript/function_libs/torch_aten/ops/core.py +++ b/onnxscript/function_libs/torch_aten/ops/core.py @@ -1643,7 +1643,7 @@ def aten_exp2(self: TFloat) -> TFloat: @torch_op("aten::expand") -def aten_expand(self: TTensor, size: IntType) -> TTensor: +def aten_expand(self: TTensor, size: TInt) -> TTensor: # expand(Tensor(a) self, SymInt[] size, *, bool implicit=False) -> Tensor(a) size = op.Cast(size, to=INT64.dtype) @@ -2235,7 +2235,7 @@ def aten_index_reduce( # FIXME(#277): Script when attributes can come before inputs @torch_op("aten::index_select", trace_only=True) -def aten_index_select(self: TTensor, dim: int, index: IntType) -> TTensor: +def aten_index_select(self: TTensor, dim: int, index: TInt) -> TTensor: # index_select(Tensor self, int dim, Tensor index) -> Tensor if op.Size(op.Shape(self)) == 0: @@ -4090,7 +4090,7 @@ def aten_renorm(self: TensorType, p: float, dim: int, maxnorm: float) -> TensorT @torch_op("aten::repeat") -def aten_repeat(self: TTensor, repeats: IntType) -> TTensor: +def aten_repeat(self: TTensor, repeats: TInt) -> TTensor: # repeat(Tensor self, SymInt[] repeats) -> Tensor if op.Size(repeats) == 0: From 678591562ebbf55db2beaae37f96267d30e96908 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 9 Jan 2023 16:30:30 +0000 Subject: [PATCH 9/9] mypy --- onnxscript/backend/onnx_backend.py | 4 ++-- onnxscript/utils.py | 21 ++++++++++++--------- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/onnxscript/backend/onnx_backend.py b/onnxscript/backend/onnx_backend.py index 1d917ffdc0..7ccc5e9a63 100644 --- a/onnxscript/backend/onnx_backend.py +++ b/onnxscript/backend/onnx_backend.py @@ -77,8 +77,8 @@ def _read_proto_from_file(full): loaded = to_list(seq) # type: ignore[assignment] except Exception: # pylint: disable=W0703 try: - loaded = onnx.load_model_from_string(serialized) - except Exception: # pragma: no cover + loaded = onnx.load_model_from_string(serialized) # type: ignore[assignment] + except Exception: raise RuntimeError( f"Unable to read {full!r}, error is {e}, " f"content is {serialized[:100]!r}." diff --git a/onnxscript/utils.py b/onnxscript/utils.py index 988016bc79..e354c204da 100644 --- a/onnxscript/utils.py +++ b/onnxscript/utils.py @@ -5,12 +5,13 @@ from __future__ import annotations import numbers -from typing import Any, Optional, Sequence +from typing import Any, Iterable, Optional, Sequence import numpy as np import onnx +import onnx.helper +import onnx.mapping from onnx import FunctionProto, ModelProto, TensorProto, ValueInfoProto -from onnx.helper import make_sequence_type_proto, make_tensor_type_proto from onnxscript import tensor @@ -82,22 +83,24 @@ def value_to_type_proto(val): if isinstance(val, (np.ndarray, tensor.Tensor)): elem_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[val.dtype] shape = val.shape - return make_tensor_type_proto(elem_type, shape) + return onnx.helper.make_tensor_type_proto(elem_type, shape) if isinstance(val, int): - return make_tensor_type_proto(TensorProto.INT32, []) + return onnx.helper.make_tensor_type_proto(TensorProto.INT32, []) if isinstance(val, (float, np.float32)): - return make_tensor_type_proto(TensorProto.FLOAT, []) + return onnx.helper.make_tensor_type_proto(TensorProto.FLOAT, []) if isinstance(val, list): if len(val) > 0: - return make_sequence_type_proto(value_to_type_proto(val[0])) + return onnx.helper.make_sequence_type_proto(value_to_type_proto(val[0])) # Edge-case. Cannot determine a suitable ONNX type for an empty list. # Should be using a typed-value instead. # Treated as a sequence of tensors of float-type. - return make_sequence_type_proto(make_tensor_type_proto(TensorProto.FLOAT, None)) + return onnx.helper.make_sequence_type_proto( + onnx.helper.make_tensor_type_proto(TensorProto.FLOAT, None) + ) if isinstance(val, numbers.Number): nparray = np.array(val) elem_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[nparray.dtype] - return make_tensor_type_proto(elem_type, []) + return onnx.helper.make_tensor_type_proto(elem_type, []) raise ValueError(f"Value of type {type(val)} is invalid as an ONNX input/output.") @@ -144,7 +147,7 @@ def make_model_from_function_proto( **(attrs or {}), ) graph = onnx.helper.make_graph([node], "node_graph", input_value_infos, output_value_infos) - model_proto_opset = function_proto.opset_import + model_proto_opset: Iterable[onnx.OperatorSetIdProto] = function_proto.opset_import if all(o.domain != function_proto.domain for o in model_proto_opset): model_proto_opset = [ *model_proto_opset,