From ba930bd9046ad0cf820f3796f3aad0428446e3b7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:46:45 +0000 Subject: [PATCH 1/4] Initial plan From 8334fcb179c74c2dd3b1c9eef3836ced0172393e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:51:00 +0000 Subject: [PATCH 2/4] Fix arange and logit behavior for torch-nightly CI (cherry-pick 6c844d3) --- .../function_libs/torch_lib/ops/core.py | 91 +++++++++++++++++-- 1 file changed, 83 insertions(+), 8 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index adf1bad4b6..231faaf83c 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -56,6 +56,9 @@ _INT64_MAX = 9223372036854775807 _INT64_MIN = -9223372036854775808 _MATH_PI = math.pi +_INT64_ARANGE_NON_INTEGRAL_USES_FLOAT_LENGTH = ( + torch.arange(3.1, dtype=torch.int64).numel() == 4 +) @torch_op("aten::_local_scalar_dense", trace_only=True) @@ -521,6 +524,56 @@ def _range_supported(dtype: int) -> bool: } +def _is_integral_dtype(dtype: int) -> bool: + return dtype in { + INT8.dtype, + INT16.dtype, + INT32.dtype, + INT64.dtype, + } + + +def _is_integral_scalar(arg: TRealUnlessFloat16OrInt8) -> bool: + if isinstance(arg, int): + return True + if isinstance(arg, float): + return False + return arg.dtype in { + INT8.dtype, + INT16.dtype, + INT32.dtype, + INT64.dtype, + } + + +def _arange_integral_dtype_with_non_integral_args( + start: TRealUnlessFloat16OrInt8, + end: TRealUnlessFloat16OrInt8, + step: TRealUnlessFloat16OrInt8, + dtype: int, +) -> TensorType: + """Implements torch.arange for integral dtypes when not all inputs are integral.""" + start_float = op.Cast(start, to=FLOAT.dtype) + end_float = op.Cast(end, to=FLOAT.dtype) + step_float = op.Cast(step, to=FLOAT.dtype) + + length = op.Cast( + op.Ceil(op.Div(op.Sub(end_float, start_float), step_float)), to=INT64.dtype + ) + index = op.Range(op.Constant(value_int=0), length, op.Constant(value_int=1)) + + if _range_supported(dtype): + index = op.Cast(index, to=dtype) + start = op.Cast(start, to=dtype) + step = op.Cast(step, to=dtype) + return op.Add(start, op.Mul(step, index)) + + start = op.Cast(start, to=INT64.dtype) + step = op.Cast(step, to=INT64.dtype) + result = op.Add(start, op.Mul(step, index)) + return op.Cast(result, to=dtype) + + def _integral_to_be_adjusted(dtype: int) -> bool: """Returns true if the dtype is special integral handled by torch.""" return dtype in { @@ -544,6 +597,12 @@ def aten_arange( zero = op.CastLike(0.0, end) one = op.CastLike(1.0, end) result = op.Range(zero, end, one) + elif ( + _is_integral_dtype(dtype) + and not _is_integral_scalar(end) + and (dtype != INT64.dtype or _INT64_ARANGE_NON_INTEGRAL_USES_FLOAT_LENGTH) + ): + result = _arange_integral_dtype_with_non_integral_args(0, end, 1, dtype) elif _range_supported(dtype): end = op.Cast(end, to=dtype) zero = op.Cast(0, to=dtype) @@ -576,6 +635,12 @@ def aten_arange_start( if dtype == -1 or dtype is None: one = op.CastLike(1.0, end) result = op.Range(start, end, one) + elif ( + _is_integral_dtype(dtype) + and not (_is_integral_scalar(start) and _is_integral_scalar(end)) + and (dtype != INT64.dtype or _INT64_ARANGE_NON_INTEGRAL_USES_FLOAT_LENGTH) + ): + result = _arange_integral_dtype_with_non_integral_args(start, end, 1, dtype) elif _range_supported(dtype): end = op.Cast(end, to=dtype) start = op.Cast(start, to=dtype) @@ -659,17 +724,26 @@ def aten_arange_start_step( end = op.Cast(end, to=FLOAT.dtype) step = op.Cast(step, to=FLOAT.dtype) result = op.Range(start, end, step) - elif _integral_to_be_adjusted(dtype): - # PyTorch arange op handles these integral types differently from INT64, - # so we have to adjust these arguments accordingly. - # https://github.com/pytorch/pytorch/blob/121cfb60c0817816fcbe2190303b7f6d05c77cf3/torch/_refs/__init__.py#L4794 - start, end, step = _adjust_args_for_arange_int_dtype(start, end, step) - result = op.Cast(op.Range(start, end, step), to=dtype) + elif ( + _is_integral_dtype(dtype) + and not ( + _is_integral_scalar(start) + and _is_integral_scalar(end) + and _is_integral_scalar(step) + ) + and (dtype != INT64.dtype or _INT64_ARANGE_NON_INTEGRAL_USES_FLOAT_LENGTH) + ): + result = _arange_integral_dtype_with_non_integral_args(start, end, step, dtype) elif dtype == INT64.dtype: end = op.Cast(end, to=dtype) start = op.Cast(start, to=dtype) step = op.Cast(step, to=dtype) result = op.Range(start, end, step) + elif _integral_to_be_adjusted(dtype): + # PyTorch arange op handles these integral types differently from INT64 + # when all arguments are integral. + start, end, step = _adjust_args_for_arange_int_dtype(start, end, step) + result = op.Cast(op.Range(start, end, step), to=dtype) else: # Cast input to float if dtype is not supported by Range, # because the input dtype may be e.g. bfloat16, @@ -5804,8 +5878,9 @@ def aten_logit(self: TFloat, eps: Optional[float] = None) -> TFloat: one_minus_eps = ir.tensor(1 - eps, dtype=self.dtype) eps = ir.tensor(eps, dtype=self.dtype) - temporary_self = op.Where(self <= one_minus_eps, self, one_minus_eps) - z = op.Where(temporary_self < eps, eps, temporary_self) + # Match torch.clamp behavior for eps > 0.5 by applying max then min. + z = op.Where(self < eps, eps, self) + z = op.Where(z <= one_minus_eps, z, one_minus_eps) return op.Log(op.Div(z, op.Sub(one, z))) From 37dec37df3e8244768352c0c4a3f48ac05b7743a Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 21 Aug 2026 10:33:33 -0700 Subject: [PATCH 3/4] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- 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 231faaf83c..39a6ceb978 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -57,7 +57,7 @@ _INT64_MIN = -9223372036854775808 _MATH_PI = math.pi _INT64_ARANGE_NON_INTEGRAL_USES_FLOAT_LENGTH = ( - torch.arange(3.1, dtype=torch.int64).numel() == 4 + torch.arange(3.1, dtype=torch.int64, device="cpu").numel() == 4 ) From d54112a9b1e5b637bb01ccc88e9a608cbf5d76ca Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:37:08 +0000 Subject: [PATCH 4/4] Compute arange torch-behavior probe lazily Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> --- .../function_libs/torch_lib/ops/core.py | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 39a6ceb978..63cbe3f7e2 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -11,6 +11,7 @@ from __future__ import annotations +import functools import math from typing import Any, Optional, Sequence, Tuple, Union @@ -56,9 +57,17 @@ _INT64_MAX = 9223372036854775807 _INT64_MIN = -9223372036854775808 _MATH_PI = math.pi -_INT64_ARANGE_NON_INTEGRAL_USES_FLOAT_LENGTH = ( - torch.arange(3.1, dtype=torch.int64, device="cpu").numel() == 4 -) + + +@functools.lru_cache(maxsize=None) +def _int64_arange_non_integral_uses_float_length() -> bool: + """Whether torch computes the arange length in float when dtype is integral. + + Newer versions of torch return 4 elements for ``torch.arange(3.1, dtype=torch.int64)``, + while older versions return 3. This is evaluated lazily to avoid running torch code + at import time. + """ + return torch.arange(3.1, dtype=torch.int64, device="cpu").numel() == 4 @torch_op("aten::_local_scalar_dense", trace_only=True) @@ -600,7 +609,7 @@ def aten_arange( elif ( _is_integral_dtype(dtype) and not _is_integral_scalar(end) - and (dtype != INT64.dtype or _INT64_ARANGE_NON_INTEGRAL_USES_FLOAT_LENGTH) + and (dtype != INT64.dtype or _int64_arange_non_integral_uses_float_length()) ): result = _arange_integral_dtype_with_non_integral_args(0, end, 1, dtype) elif _range_supported(dtype): @@ -638,7 +647,7 @@ def aten_arange_start( elif ( _is_integral_dtype(dtype) and not (_is_integral_scalar(start) and _is_integral_scalar(end)) - and (dtype != INT64.dtype or _INT64_ARANGE_NON_INTEGRAL_USES_FLOAT_LENGTH) + and (dtype != INT64.dtype or _int64_arange_non_integral_uses_float_length()) ): result = _arange_integral_dtype_with_non_integral_args(start, end, 1, dtype) elif _range_supported(dtype): @@ -731,7 +740,7 @@ def aten_arange_start_step( and _is_integral_scalar(end) and _is_integral_scalar(step) ) - and (dtype != INT64.dtype or _INT64_ARANGE_NON_INTEGRAL_USES_FLOAT_LENGTH) + and (dtype != INT64.dtype or _int64_arange_non_integral_uses_float_length()) ): result = _arange_integral_dtype_with_non_integral_args(start, end, step, dtype) elif dtype == INT64.dtype: