From a6712b7fb11a413217e05b3301bbabaf5c477206 Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Thu, 23 Oct 2025 20:00:58 +0000 Subject: [PATCH 1/3] support match trunc --- onnxscript/function_libs/torch_lib/ops/core.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 36f2a70f8c..c1af1247f0 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -8533,6 +8533,14 @@ def aten_trunc(self: TFloat) -> TFloat: return op.Floor(op.Abs(self)) * op.Sign(self) +@torch_op("math::trunc", trace_only=True) +def python_math_trunc(self: TFloat) -> TInt: + """trunc(Tensor self) -> Tensor""" + # Reference https://github.com/onnx/onnx/issues/4588#issuecomment-2658170591 + result = op.Floor(op.Abs(self)) * op.Sign(self) + return op.Cast(result, to=INT64.dtype) + + @torch_op("aten::type_as", trace_only=True) def aten_type_as(self: TTensor, other: TTensor2) -> TTensor2: """type_as(Tensor self, Tensor other) -> Tensor""" From 7d30b1fb1d9cd13aba095f85295781937abb7ec6 Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Mon, 27 Oct 2025 18:18:17 +0000 Subject: [PATCH 2/3] only use cast --- 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 c1af1247f0..c0a8db3847 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -8536,9 +8536,7 @@ def aten_trunc(self: TFloat) -> TFloat: @torch_op("math::trunc", trace_only=True) def python_math_trunc(self: TFloat) -> TInt: """trunc(Tensor self) -> Tensor""" - # Reference https://github.com/onnx/onnx/issues/4588#issuecomment-2658170591 - result = op.Floor(op.Abs(self)) * op.Sign(self) - return op.Cast(result, to=INT64.dtype) + return op.Cast(self, to=INT64.dtype) @torch_op("aten::type_as", trace_only=True) From 6a084dd9823117938888483c537d1fd3643ec47a Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Mon, 27 Oct 2025 18:19:30 +0000 Subject: [PATCH 3/3] add note --- onnxscript/function_libs/torch_lib/ops/core.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index c0a8db3847..eb276a239c 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -8536,6 +8536,8 @@ def aten_trunc(self: TFloat) -> TFloat: @torch_op("math::trunc", trace_only=True) def python_math_trunc(self: TFloat) -> TInt: """trunc(Tensor self) -> Tensor""" + # NOTE: This is used in SymInt/SymBool/SymFloat context, so + # we don't expect overflow to happen here. return op.Cast(self, to=INT64.dtype)