onnxscript 0.7.2.dev20260825__tar.gz → 0.7.2.dev20260903__tar.gz

{onnxscript-0.7.2.dev20260825/onnxscript.egg-info → onnxscript-0.7.2.dev20260903}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxscript
-Version: 0.7.2.dev20260825
+Version: 0.7.2.dev20260903
 Summary: Naturally author ONNX functions and models using a subset of Python
 Author-email: Microsoft Corporation <onnx@microsoft.com>
 License: MIT License
@@ -27,7 +27,7 @@ License: MIT License
         
 Project-URL: Homepage, https://microsoft.github.io/onnxscript/
 Project-URL: Repository, https://github.com/microsoft/onnxscript
-Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/f78d33b568ec88ea78c3e5fbac89f703963f0e29
+Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/b4c10fab5b65d7f15bbaadae89512ba4fba74745
 Classifier: Development Status :: 4 - Beta
 Classifier: Environment :: Console
 Classifier: Intended Audience :: Developers

{onnxscript-0.7.2.dev20260825 → onnxscript-0.7.2.dev20260903}/onnxscript/_framework_apis/torch_2_5.py

@@ -59,9 +59,13 @@ def convert_version(model: ir.Model, target_version: int) -> ir.Model:
 
 
 def check_model(model: ir.Model) -> None:
-    """Check the model."""
+    """No-op retained for API compatibility.
 
-    del model  # Unused yet
+    This intentionally performs no validation. Running the ONNX checker here was
+    dropped because it can report false positives and adds overhead to export.
+    """
+
+    del model  # Intentionally unused: this function performs no validation.
 
 
 def save_model_with_external_data(

{onnxscript-0.7.2.dev20260825 → onnxscript-0.7.2.dev20260903}/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
 
@@ -58,6 +59,17 @@ _INT64_MIN = -9223372036854775808
 _MATH_PI = math.pi
 
 
+@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)
 def aten__local_scalar_dense(self: TensorType) -> TensorType:
     """_local_scalar_dense(Tensor self) -> Scalar"""
@@ -521,6 +533,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 +606,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 +644,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 +733,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,
@@ -3163,7 +3246,10 @@ def aten_div_mode(self: TReal, other: TReal, rounding_mode: Optional[str] = None
     if rounding_mode == "trunc":
         # Rounds the results of the division towards zero.
         # Equivalent to C-style integer division
-        return aten_trunc(op.Div(self, other))
+        quotient = op.Div(self, other)
+        if self.dtype == ir.DataType.FLOAT16:
+            quotient = op.Cast(quotient, to=FLOAT16.dtype)
+        return aten_trunc(quotient)
     if rounding_mode == "floor":
         return op.Floor(op.Div(self, other))
 
@@ -5445,7 +5531,7 @@ def aten_is_vulkan_available() -> bool:
     raise NotImplementedError()
 
 
-@torch_op("aten::isclose")
+@torch_op("aten::isclose", trace_only=True)
 def aten_isclose(
     self: TReal,
     other: TReal,
@@ -5455,12 +5541,37 @@ def aten_isclose(
 ) -> BOOL:
     """isclose(Tensor self, Tensor other, float rtol=1e-05, float atol=1e-08, bool equal_nan=False) -> Tensor"""
 
-    # FIXME: check equal_nan when self and other are all NaN
-    # |input - other| <= atol + rtol x |other|
-    left_part = op.Abs(op.Sub(self, other))
-    right_part = op.Add(atol, op.Mul(rtol, op.Abs(other)))
-    result = op.LessOrEqual(left_part, right_part)
-    return result
+    # torch builds isclose out of three terms in aten/src/ATen/native/TensorCompare.cpp:
+    # exact equality, NaN against NaN when equal_nan is set, and the tolerance band
+    # restricted to the elements whose error is finite.
+    result = op.Equal(self, other)
+
+    # |self - other| <= atol + rtol x |other|
+    actual_error = op.Abs(op.Sub(self, other))
+    allowed_error = op.Add(
+        op.CastLike(atol, other), op.Mul(op.CastLike(rtol, other), op.Abs(other))
+    )
+
+    if self.dtype.is_floating_point():
+        # Comparing the two sides of the band by subtraction carries the finiteness
+        # restriction on its own, so no IsInf term is needed. An infinity on either side
+        # puts a NaN into the subtraction whichever way it arrives: the error and the
+        # allowance are both infinite and cancel, or the allowance is already NaN because
+        # rtol is zero and zero times infinity is NaN. Every comparison against NaN is
+        # false, which leaves the equality term above as the only way an infinity is
+        # reported close. On finite values the sign of the difference and the direct
+        # comparison agree exactly, so the band is unchanged.
+        zero = op.Constant(value=ir.tensor(0, dtype=self.dtype))
+        within_tolerance = op.LessOrEqual(op.Sub(actual_error, allowed_error), zero)
+        if equal_nan:
+            result = op.Or(result, op.And(op.IsNaN(self), op.IsNaN(other)))
+    else:
+        # Integers are never infinite and never NaN, so they need no finiteness term, and
+        # IsNaN does not accept them. The subtraction is left off here because it can
+        # overflow a narrow integer type where the direct comparison cannot.
+        within_tolerance = op.LessOrEqual(actual_error, allowed_error)
+
+    return op.Or(result, within_tolerance)
 
 
 @torch_op("aten::isfinite")
@@ -5852,8 +5963,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)))
 

{onnxscript-0.7.2.dev20260825 → onnxscript-0.7.2.dev20260903}/onnxscript/rewriter/rules/common/_fuse_batchnorm.py

@@ -69,6 +69,14 @@ class _FuseBatchNormBase(RewriteRuleClassBase, ABC):
             self._scale_weights(weights, scale_factor, inbound_node.attributes)
         )
 
+        # Gemm optionally scales its bias input C by the "beta" attribute (default 1.0,
+        # https://onnx.ai/onnx/operators/onnx__Gemm.html#attributes); Conv/ConvTranspose have
+        # no such attribute. Fold that scaling into the fused bias now, and drop "beta" from
+        # the re-emitted node's attributes so it is not applied a second time there.
+        new_attributes = dict(inbound_node.attributes)
+        gemm_beta_attr = new_attributes.pop("beta", None)
+        gemm_beta = gemm_beta_attr.as_float() if gemm_beta_attr is not None else 1.0
+
         # Update bias
         if len(inbound_node.inputs) > 2:
             original_bias = inbound_node.inputs[2].const_value.numpy()
@@ -79,14 +87,14 @@ class _FuseBatchNormBase(RewriteRuleClassBase, ABC):
             # to avoid name collision on initializer creation when there are multiple patterns
             # sharing the same parent nodes.
             bias_name = inbound_node.inputs[1].name + "_bias"
-        fused_bias = ir.tensor((original_bias - input_mean) * scale_factor + beta)
+        fused_bias = ir.tensor((gemm_beta * original_bias - input_mean) * scale_factor + beta)
 
         return op.op(
             self.op_type,
             x,
             op.initializer(fused_weights, name=inbound_node.inputs[1].name),
             op.initializer(fused_bias, name=bias_name),
-            **inbound_node.attributes,
+            **new_attributes,
         )
 
     def check(self, context, x, inbound_out: ir.Value, batchnorm_out: ir.Value) -> MatchResult:

{onnxscript-0.7.2.dev20260825 → onnxscript-0.7.2.dev20260903}/onnxscript/version_converter/_version_converter.py

@@ -160,11 +160,14 @@ def dft_19_20(node: ir.Node, op):
     dft_length = node.inputs[1] if len(node.inputs) > 1 else None
     inverse = _get_int_attribute(node, "inverse", 0)
     onesided = _get_int_attribute(node, "onesided", 0)
-    axis = _get_int_attribute(node, "axis", None)
-    if axis is not None:
-        axis_value = op.Constant(value_int=axis)
-        return op.DFT(input, dft_length, axis_value, inverse=inverse, onesided=onesided)
-    return None
+    # In opset 19 `axis` is an attribute defaulting to 1; in opset 20 it became an
+    # input defaulting to -2. An omitted attribute must therefore be materialized
+    # explicitly, or the converted node silently transforms a different axis.
+    axis = _get_int_attribute(node, "axis", 1)
+    if axis is None:
+        return None
+    axis_value = op.Constant(value_int=axis)
+    return op.DFT(input, dft_length, axis_value, inverse=inverse, onesided=onesided)
 
 
 @register("GridSample", node_version=19, up_conversion=True)

{onnxscript-0.7.2.dev20260825 → onnxscript-0.7.2.dev20260903/onnxscript.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxscript
-Version: 0.7.2.dev20260825
+Version: 0.7.2.dev20260903
 Summary: Naturally author ONNX functions and models using a subset of Python
 Author-email: Microsoft Corporation <onnx@microsoft.com>
 License: MIT License
@@ -27,7 +27,7 @@ License: MIT License
         
 Project-URL: Homepage, https://microsoft.github.io/onnxscript/
 Project-URL: Repository, https://github.com/microsoft/onnxscript
-Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/f78d33b568ec88ea78c3e5fbac89f703963f0e29
+Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/b4c10fab5b65d7f15bbaadae89512ba4fba74745
 Classifier: Development Status :: 4 - Beta
 Classifier: Environment :: Console
 Classifier: Intended Audience :: Developers