From 6f3f5eca8a67a926a5656aa3dc13748177d9879b Mon Sep 17 00:00:00 2001 From: xadupre Date: Tue, 21 Jan 2025 15:40:38 +0100 Subject: [PATCH 01/11] Fix floordiv --- onnxscript/function_libs/torch_lib/ops/core.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index a1793858e9..8655661df0 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -2759,6 +2759,20 @@ def aten_div(self: TFloat, other: TFloat) -> TFloat: return op.Div(self, other) +@torch_op( + ( + "aten::true_divide.Tensor", + "aten::true_divide.Scalar", + "_operator::truediv", + ) +) +def aten_div_int(self: TInt, other: TInt) -> TFloat: + """div.Tensor(Tensor self, Tensor other) -> Tensor""" + + # Int inputs will be promoted to float by PyTorch + return op.Div(op.Cast(self, to=FLOAT.dtype), op.Cast(other, to=FLOAT.dtype)) + + @torch_op( ( "aten::div.Tensor", @@ -3605,12 +3619,12 @@ def aten_floor_divide(self: TFloat, other: TFloat) -> TFloat: @torch_op(("aten::floor_divide", "_operator::floordiv"), traceable=True) -def aten_floor_divide_int(self: TInt, other: TInt) -> TInt: +def aten_floor_divide_int(self: TInt, other: TInt) -> FLOAT: """floor_divide(Tensor self, Tensor other) -> Tensor""" # We implement floor_divide only for positive inputs (using integer division) # because that is the usual intended case and is the most efficient. - return op.Div(self, other) + return op.Div(op.Cast(self, to=FLOAT.dtype), op.Cast(other, to=FLOAT.dtype)) def aten_fmax(self: TensorType, other: TensorType) -> TensorType: From 42cde5e35f8a1d5a8ccc865df44c671cd0493d5c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Tue, 21 Jan 2025 17:56:49 +0100 Subject: [PATCH 02/11] Update onnxscript/function_libs/torch_lib/ops/core.py Co-authored-by: Justin Chu --- onnxscript/function_libs/torch_lib/ops/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 8655661df0..af4431e129 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -2766,7 +2766,7 @@ def aten_div(self: TFloat, other: TFloat) -> TFloat: "_operator::truediv", ) ) -def aten_div_int(self: TInt, other: TInt) -> TFloat: +def operator_truediv_int(self: TInt, other: TInt) -> TFloat: """div.Tensor(Tensor self, Tensor other) -> Tensor""" # Int inputs will be promoted to float by PyTorch From e999fbaf2663e43ff4c07765f0a59dd0a08bb72f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Tue, 21 Jan 2025 17:56:56 +0100 Subject: [PATCH 03/11] Update onnxscript/function_libs/torch_lib/ops/core.py Co-authored-by: Justin Chu --- onnxscript/function_libs/torch_lib/ops/core.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index af4431e129..8dec3c41d3 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -2760,11 +2760,7 @@ def aten_div(self: TFloat, other: TFloat) -> TFloat: @torch_op( - ( - "aten::true_divide.Tensor", - "aten::true_divide.Scalar", - "_operator::truediv", - ) + "_operator::truediv" ) def operator_truediv_int(self: TInt, other: TInt) -> TFloat: """div.Tensor(Tensor self, Tensor other) -> Tensor""" From b5f6df071a1839dbc4c4d29ff3e0664d05b19b0e Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jan 2025 09:02:44 -0800 Subject: [PATCH 04/11] Update onnxscript/function_libs/torch_lib/ops/core.py --- onnxscript/function_libs/torch_lib/ops/core.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 8dec3c41d3..e3c89b3381 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -2759,9 +2759,7 @@ def aten_div(self: TFloat, other: TFloat) -> TFloat: return op.Div(self, other) -@torch_op( - "_operator::truediv" -) +@torch_op("_operator::truediv", traceable=True) def operator_truediv_int(self: TInt, other: TInt) -> TFloat: """div.Tensor(Tensor self, Tensor other) -> Tensor""" From 78975f055e794b52b55f56214ef8ffc929f10af4 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jan 2025 09:10:32 -0800 Subject: [PATCH 05/11] Update onnxscript/function_libs/torch_lib/ops/core.py --- onnxscript/function_libs/torch_lib/ops/core.py | 1 - 1 file changed, 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index e3c89b3381..a44ade2ca2 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -2763,7 +2763,6 @@ def aten_div(self: TFloat, other: TFloat) -> TFloat: def operator_truediv_int(self: TInt, other: TInt) -> TFloat: """div.Tensor(Tensor self, Tensor other) -> Tensor""" - # Int inputs will be promoted to float by PyTorch return op.Div(op.Cast(self, to=FLOAT.dtype), op.Cast(other, to=FLOAT.dtype)) From 9deecf4c6159e45d651344bc4e305547ee7c2ded Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jan 2025 09:55:22 -0800 Subject: [PATCH 06/11] Refactor torch_op decorator usage and functions --- onnxscript/function_libs/torch_lib/ops/core.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index a44ade2ca2..4e97f65d72 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -2774,7 +2774,6 @@ def operator_truediv_int(self: TInt, other: TInt) -> TFloat: "aten::divide.Scalar", "aten::true_divide.Tensor", "aten::true_divide.Scalar", - "_operator::truediv", ), complex=True, ) @@ -3604,20 +3603,18 @@ def python_math_floor(self: TFloat) -> TInt: return op.Cast(floor, to=INT64.dtype) -@torch_op(("aten::floor_divide", "_operator::floordiv"), traceable=True) +@torch_op("aten::floor_divide", traceable=True) def aten_floor_divide(self: TFloat, other: TFloat) -> TFloat: """floor_divide(Tensor self, Tensor other) -> Tensor""" return op.Floor(op.Div(self, other)) -@torch_op(("aten::floor_divide", "_operator::floordiv"), traceable=True) -def aten_floor_divide_int(self: TInt, other: TInt) -> FLOAT: - """floor_divide(Tensor self, Tensor other) -> Tensor""" - +@torch_op("_operator::floordiv", traceable=True) +def operator_floordiv(self: TInt, other: TInt) -> FLOAT: # We implement floor_divide only for positive inputs (using integer division) # because that is the usual intended case and is the most efficient. - return op.Div(op.Cast(self, to=FLOAT.dtype), op.Cast(other, to=FLOAT.dtype)) + return op.Div(self, other) def aten_fmax(self: TensorType, other: TensorType) -> TensorType: From 75edb4d4f09a6099bf0a9be231af2a68b9d71dbb Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jan 2025 09:56:18 -0800 Subject: [PATCH 07/11] Fix return type of operator_floordiv function --- onnxscript/function_libs/torch_lib/ops/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 4e97f65d72..64d2461718 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -3611,7 +3611,7 @@ def aten_floor_divide(self: TFloat, other: TFloat) -> TFloat: @torch_op("_operator::floordiv", traceable=True) -def operator_floordiv(self: TInt, other: TInt) -> FLOAT: +def operator_floordiv(self: TInt, other: TInt) -> TInt: # We implement floor_divide only for positive inputs (using integer division) # because that is the usual intended case and is the most efficient. return op.Div(self, other) From 1ddc57b2bd6809414fabcc5741e484bc07e9153f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jan 2025 10:00:35 -0800 Subject: [PATCH 08/11] Remove `_operator` references from torch operations --- onnxscript/function_libs/torch_lib/ops/core.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 64d2461718..ad9b7856c6 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -170,7 +170,7 @@ def aten_add(self: TReal, other: TReal, alpha: float = 1.0) -> TReal: @torch_op( - ("aten::add.Tensor", "aten::add.Scalar", "_operator::add"), trace_only=True, complex=True + ("aten::add.Tensor", "aten::add.Scalar"), trace_only=True, complex=True ) def aten_add_complex(self: TReal, other: TReal, alpha: float = 1.0) -> TReal: """add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor""" @@ -2749,7 +2749,6 @@ def aten_dist(self: TensorType, other: TensorType, p: float = 2.0) -> TensorType "aten::divide.Scalar", "aten::true_divide.Tensor", "aten::true_divide.Scalar", - "_operator::truediv", ) ) def aten_div(self: TFloat, other: TFloat) -> TFloat: @@ -2760,7 +2759,7 @@ def aten_div(self: TFloat, other: TFloat) -> TFloat: @torch_op("_operator::truediv", traceable=True) -def operator_truediv_int(self: TInt, other: TInt) -> TFloat: +def operator_truediv(self: TensorType, other: TensorType) -> FLOAT: """div.Tensor(Tensor self, Tensor other) -> Tensor""" return op.Div(op.Cast(self, to=FLOAT.dtype), op.Cast(other, to=FLOAT.dtype)) @@ -3611,7 +3610,7 @@ def aten_floor_divide(self: TFloat, other: TFloat) -> TFloat: @torch_op("_operator::floordiv", traceable=True) -def operator_floordiv(self: TInt, other: TInt) -> TInt: +def operator_floordiv(self: INT64, other: INT64) -> INT64: # We implement floor_divide only for positive inputs (using integer division) # because that is the usual intended case and is the most efficient. return op.Div(self, other) @@ -4944,7 +4943,6 @@ def aten_logical_not(self: BOOL) -> BOOL: "aten::bitwise_or.Scalar_Tensor", "aten::add.Tensor", "aten::add.Scalar", - "_operator::add", ), traceable=True, ) @@ -5662,7 +5660,7 @@ def aten_mul(self: TReal, other: TReal) -> TReal: @torch_op( - ("aten::mul", "aten::mul.Tensor", "_operator::mul", "aten::multiply.Tensor"), + ("aten::mul", "aten::mul.Tensor", "aten::multiply.Tensor"), traceable=True, ) def aten_mul_bool(self: BOOL, other: BOOL) -> BOOL: @@ -5675,7 +5673,7 @@ def aten_mul_bool(self: BOOL, other: BOOL) -> BOOL: @torch_op( - ("aten::mul", "aten::mul.Tensor", "_operator::mul", "aten::multiply.Tensor"), + ("aten::mul", "aten::mul.Tensor", "aten::multiply.Tensor"), traceable=True, complex=True, ) @@ -8048,7 +8046,6 @@ def aten_sub(self: TReal, other: TReal, alpha: float = 1.0) -> TReal: "aten::sub.Scalar", "aten::subtract.Tensor", "aten::subtract.Scalar", - "_operator::sub", ), trace_only=True, complex=True, From d29deec4fd88d9d05a6cce046de2b7a88784c8a2 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jan 2025 10:05:58 -0800 Subject: [PATCH 09/11] Update onnxscript/function_libs/torch_lib/ops/core.py --- onnxscript/function_libs/torch_lib/ops/core.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index ad9b7856c6..d053f45519 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -2760,8 +2760,6 @@ def aten_div(self: TFloat, other: TFloat) -> TFloat: @torch_op("_operator::truediv", traceable=True) def operator_truediv(self: TensorType, other: TensorType) -> FLOAT: - """div.Tensor(Tensor self, Tensor other) -> Tensor""" - return op.Div(op.Cast(self, to=FLOAT.dtype), op.Cast(other, to=FLOAT.dtype)) From 6c7e9fc22c6627902d9d753ccc9fcdc9f29b0f95 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jan 2025 10:26:46 -0800 Subject: [PATCH 10/11] Remove redundant decorator line breaks --- onnxscript/function_libs/torch_lib/ops/core.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index d053f45519..58b2ae3211 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -169,9 +169,7 @@ def aten_add(self: TReal, other: TReal, alpha: float = 1.0) -> TReal: return op.Add(self, other) -@torch_op( - ("aten::add.Tensor", "aten::add.Scalar"), trace_only=True, complex=True -) +@torch_op(("aten::add.Tensor", "aten::add.Scalar"), trace_only=True, complex=True) def aten_add_complex(self: TReal, other: TReal, alpha: float = 1.0) -> TReal: """add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor""" From 7307162b068e2129a3e9b9c2311d36ac57843a92 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jan 2025 10:28:01 -0800 Subject: [PATCH 11/11] Remove unused TorchLibOpInfo for floor_divide. --- tests/function_libs/torch_lib/ops_test_data.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index ee86327362..91e10b4097 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/tests/function_libs/torch_lib/ops_test_data.py @@ -829,7 +829,6 @@ def _where_input_wrangler( test_class_name="TestOutputConsistencyEager", reason="fixme: off-by-one issue due to numerical precision. https://github.com/microsoft/onnxscript/issues/989", ), - TorchLibOpInfo("ops.aten.floor_divide.int", core_ops.aten_floor_divide_int), TorchLibOpInfo("fmod", core_ops.aten_fmod), TorchLibOpInfo("frac", core_ops.aten_frac), TorchLibOpInfo("full", core_ops.aten_full),