onnxscript 0.7.2.dev20260826__tar.gz → 0.7.2.dev20260904__tar.gz

{onnxscript-0.7.2.dev20260826/onnxscript.egg-info → onnxscript-0.7.2.dev20260904}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxscript
-Version: 0.7.2.dev20260826
+Version: 0.7.2.dev20260904
 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/3ba2bf74970122d2a8b043765128f34cea69be8f
+Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/9709f8a0777037c4c0a6320a64593f342aef7183
 Classifier: Development Status :: 4 - Beta
 Classifier: Environment :: Console
 Classifier: Intended Audience :: Developers

{onnxscript-0.7.2.dev20260826 → onnxscript-0.7.2.dev20260904}/onnxscript/function_libs/torch_lib/ops/core.py

@@ -3246,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))
 
@@ -4670,6 +4673,27 @@ def aten_grid_sampler_3d_backward(
     raise NotImplementedError()
 
 
+@torch_op("aten::_grouped_mm", trace_only=True)
+def aten_grouped_mm(
+    self: TFloat,
+    mat2: TFloat,
+    offs: Optional[TInt] = None,
+    bias: Optional[TFloat] = None,
+    out_dtype: int = -1,
+) -> TFloat:
+    """_grouped_mm(Tensor self, Tensor mat2, *, Tensor? offs=None, Tensor? bias=None, int? out_dtype=None) -> Tensor"""
+
+    if offs is not None:
+        raise NotImplementedError("Grouped matmul with offsets (ragged/MoE) is not supported.")
+
+    res = op.MatMul(self, mat2)
+    if bias is not None:
+        res = op.Add(res, bias)
+    if out_dtype is not None and out_dtype != -1:
+        res = op.Cast(res, to=out_dtype)
+    return res
+
+
 def aten_gru_cell(
     input: TensorType,
     hx: TensorType,
@@ -5528,7 +5552,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,
@@ -5538,12 +5562,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")

{onnxscript-0.7.2.dev20260826 → onnxscript-0.7.2.dev20260904}/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.dev20260826 → onnxscript-0.7.2.dev20260904}/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.dev20260826 → onnxscript-0.7.2.dev20260904/onnxscript.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxscript
-Version: 0.7.2.dev20260826
+Version: 0.7.2.dev20260904
 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/3ba2bf74970122d2a8b043765128f34cea69be8f
+Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/9709f8a0777037c4c0a6320a64593f342aef7183
 Classifier: Development Status :: 4 - Beta
 Classifier: Environment :: Console
 Classifier: Intended Audience :: Developers