onnxscript 0.7.2.dev20260714__tar.gz → 0.7.2.dev20260731__tar.gz

{onnxscript-0.7.2.dev20260714/onnxscript.egg-info → onnxscript-0.7.2.dev20260731}/PKG-INFO

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

{onnxscript-0.7.2.dev20260714 → onnxscript-0.7.2.dev20260731}/onnxscript/function_libs/torch_lib/ops/core.py

@@ -1254,12 +1254,41 @@ def aten_binary_cross_entropy_with_logits(
     raise NotImplementedError()
 
 
+@torch_op("aten::bincount", trace_only=True)
 def aten_bincount(
-    self: TensorType, weights: Optional[TensorType] = None, minlength: int = 0
+    self: IntType, weights: Optional[TensorType] = None, minlength: int = 0
 ) -> TensorType:
-    """bincount(Tensor self, Tensor? weights=None, int minlength=0) -> Tensor"""
+    """bincount(Tensor self, Tensor? weights=None, int minlength=0) -> Tensor
 
-    raise NotImplementedError()
+    ``weights`` is not supported. Negative inputs are rejected by torch and are not
+    handled here (ONNX integer ops would wrap them around).
+    """
+    if weights is not None:
+        raise NotImplementedError("aten::bincount with weights is not supported.")
+
+    self = op.Cast(self, to=INT64.dtype)
+    axis_0 = op.Constant(value_ints=[0])
+    # Append a 0 so ReduceMax is defined even when ``self`` is empty. It only sizes the
+    # output and never contributes to the counts (the scatter below uses ``self``).
+    data_max = op.Unsqueeze(
+        op.ReduceMax(op.Concat(self, op.Constant(value_ints=[0]), axis=0), keepdims=0),
+        axis_0,
+    )
+    # An empty input yields depth 0, so the output is empty unless ``minlength`` applies.
+    non_empty = op.Unsqueeze(
+        op.Cast(op.Greater(op.Size(self), op.Constant(value_int=0)), to=INT64.dtype),
+        axis_0,
+    )
+    depth = op.Mul(op.Add(data_max, op.Constant(value_ints=[1])), non_empty)
+    if minlength > 0:
+        depth = op.Max(depth, op.Constant(value_ints=[minlength]))
+
+    # Scatter-add 1 for each value into a zero vector of length ``depth``. This uses
+    # O(N + depth) memory instead of the dense O(N * depth) one-hot, and behaves
+    # correctly for empty inputs.
+    zeros = op.Expand(op.Constant(value_int=0), depth)
+    ones = op.Expand(op.Constant(value_int=1), op.Shape(self))
+    return op.ScatterElements(zeros, self, ones, axis=0, reduction="add")
 
 
 def aten_binomial(
@@ -2276,6 +2305,29 @@ def aten_convolution(
     return result
 
 
+def _conv_kernel_shape(weight: TFloat, complex: bool = False) -> Optional[Sequence[int]]:
+    """Return the spatial (kernel) shape of the convolution weight if statically known.
+
+    The kernel_shape attribute of ONNX Conv/ConvTranspose corresponds to the spatial
+    dimensions of the weight tensor, i.e. all dimensions except the leading two
+    (output and input channels). When ``complex`` is True, the weight has a trailing
+    dimension of size 2 (real/imaginary parts) which is also excluded. Returns None
+    when the spatial shape is not fully static, in which case the attribute is omitted
+    and inferred by ONNX.
+    """
+    shape = weight.shape
+    if shape is None:
+        return None
+    # Exclude the leading output/input channel dims, and for complex weights the
+    # trailing real/imag dim as well.
+    kernel_shape = shape[2:-1] if complex else shape[2:]
+    if len(kernel_shape) == 0:
+        return None
+    if any(not isinstance(dim, int) for dim in kernel_shape):
+        return None
+    return list(kernel_shape)
+
+
 def _aten_convolution_onnx(
     input: TFloat,
     weight: TFloat,
@@ -2299,6 +2351,11 @@ def _aten_convolution_onnx(
     if no_batch:
         input = op.Unsqueeze(input, op.Constant(value_ints=[0]))
 
+    # kernel_shape is the spatial shape of the weight tensor (excluding the
+    # output-channel and input-channel dimensions). It is optional in the ONNX
+    # spec but is set explicitly here to match the ONNX Conv specification.
+    kernel_shape = _conv_kernel_shape(weight)
+
     if transposed:
         result = op.ConvTranspose(
             input,
@@ -2309,6 +2366,7 @@ def _aten_convolution_onnx(
             group=groups,
             dilations=dilations,
             output_padding=output_padding,
+            kernel_shape=kernel_shape,
         )
     else:
         result = op.Conv(
@@ -2319,6 +2377,7 @@ def _aten_convolution_onnx(
             pads=pads,
             group=groups,
             dilations=dilations,
+            kernel_shape=kernel_shape,
         )
 
     if no_batch:
@@ -2358,6 +2417,11 @@ def _aten_convolution_complex_onnx(
     bias_imag = op.Gather(bias, 1, axis=-1)
     bias_zero = op.Expand(op.CastLike(0.0, weight), op.Shape(bias_real))
 
+    # The complex weight has a trailing dimension of size 2 (real/imaginary parts),
+    # so the spatial (kernel) shape excludes both leading channel dimensions and the
+    # trailing complex dimension.
+    kernel_shape = _conv_kernel_shape(weight, complex=True)
+
     if transposed:
         result_real = op.Sub(
             op.ConvTranspose(
@@ -2369,6 +2433,7 @@ def _aten_convolution_complex_onnx(
                 group=groups,
                 dilations=dilations,
                 output_padding=output_padding,
+                kernel_shape=kernel_shape,
             ),
             op.ConvTranspose(
                 input_imag,
@@ -2379,6 +2444,7 @@ def _aten_convolution_complex_onnx(
                 group=groups,
                 dilations=dilations,
                 output_padding=output_padding,
+                kernel_shape=kernel_shape,
             ),
         )
         result_imag = op.Add(
@@ -2391,6 +2457,7 @@ def _aten_convolution_complex_onnx(
                 group=groups,
                 dilations=dilations,
                 output_padding=output_padding,
+                kernel_shape=kernel_shape,
             ),
             op.ConvTranspose(
                 input_imag,
@@ -2401,6 +2468,7 @@ def _aten_convolution_complex_onnx(
                 group=groups,
                 dilations=dilations,
                 output_padding=output_padding,
+                kernel_shape=kernel_shape,
             ),
         )
     else:
@@ -2413,6 +2481,7 @@ def _aten_convolution_complex_onnx(
                 pads=pads,
                 group=groups,
                 dilations=dilations,
+                kernel_shape=kernel_shape,
             ),
             op.Conv(
                 input_imag,
@@ -2422,6 +2491,7 @@ def _aten_convolution_complex_onnx(
                 pads=pads,
                 group=groups,
                 dilations=dilations,
+                kernel_shape=kernel_shape,
             ),
         )
 
@@ -2434,6 +2504,7 @@ def _aten_convolution_complex_onnx(
                 pads=pads,
                 group=groups,
                 dilations=dilations,
+                kernel_shape=kernel_shape,
             ),
             op.Conv(
                 input_imag,
@@ -2443,6 +2514,7 @@ def _aten_convolution_complex_onnx(
                 pads=pads,
                 group=groups,
                 dilations=dilations,
+                kernel_shape=kernel_shape,
             ),
         )
 
@@ -4983,8 +5055,14 @@ def aten_index_put(
         # will invalidate equality-based check.
         first_shape = indices[advanced_indices[0]].shape
 
-        def same_shape(other_shape: ir.Shape) -> bool:
-            return (not any(d is None for d in other_shape)) and other_shape == first_shape
+        def same_shape(other_shape: Optional[ir.Shape]) -> bool:
+            return (
+                first_shape is not None
+                and other_shape is not None
+                and not any(d is None for d in first_shape)
+                and not any(d is None for d in other_shape)
+                and other_shape == first_shape
+            )
 
         all_same_shape = all(same_shape(indices[i].shape) for i in advanced_indices)
         if not all_same_shape:
@@ -5078,24 +5156,70 @@ def aten_index_put(
 
 def _aten_index_put_bool(
     self: TReal,
-    indices: Sequence[BOOL],
+    indices: Sequence[Optional[Union[INT64, BOOL]]],
     values: TReal,
     accumulate: bool = False,
 ) -> TReal:
     """index_put(Tensor self, Tensor?[] indices, Tensor values, bool accumulate=False) -> Tensor"""
 
-    # TODO: Support indices with more than 1 elements
-    index = indices[0]
-    # accumulate should be always False, True does not make sense but an assert would be great
-    # Reshape indices so it can be properly broadcasted
+    bool_mask = indices[0]
+    if len(indices) > 1:
+        if any(index is None for index in indices):
+            raise NotImplementedError(
+                "Boolean index_put with multiple indices does not support None indices."
+            )
+
+        advanced_indices = []
+        selected_positions = []
+        minus_one = op.Constant(value_ints=[-1])
+        for index in indices:
+            if index.dtype != BOOL.dtype or len(index.shape) != 1:
+                raise NotImplementedError(
+                    "Boolean index_put with multiple indices supports only 1-D boolean masks."
+                )
+            positions = op.Reshape(op.Transpose(op.NonZero(index), perm=[1, 0]), minus_one)
+            selected_positions.append(positions)
+            advanced_indices.append(op.Unsqueeze(positions, minus_one))
+        onnx_index = op.Concat(*advanced_indices, axis=-1)
+        target_shape = op.Concat(
+            op.Shape(selected_positions[0]),
+            op.Slice(op.Shape(self), starts=[len(indices)], ends=[len(self.shape)], axes=[0]),
+            axis=0,
+        )
+        expanded_values = op.Expand(values, target_shape)
+        return op.ScatterND(
+            self, onnx_index, expanded_values, reduction="add" if accumulate else None
+        )
+
+    if bool_mask is None or bool_mask.dtype != BOOL.dtype:
+        raise NotImplementedError(
+            "Boolean index_put expects a boolean mask as the first index."
+        )
+
+    neg_1 = op.Constant(value_ints=[-1])
     self_rank = len(self.shape)
-    index_rank = len(index.shape)
-    if self_rank > index_rank:
-        index_shape = op.Shape(index)
-        padding = op.Constant(value_ints=[1 for _ in range(self_rank - index_rank)])
-        padded_shape = op.Concat(index_shape, padding, axis=0)
-        index = op.Reshape(index, padded_shape)
-    return op.Where(index, values, self)
+    mask_rank = len(bool_mask.shape)
+
+    # Expand a lower-rank mask (e.g. a row mask) across the trailing dimensions of self
+    # so it selects whole slices, then collect the coordinates of every selected element.
+    # NonZero returns them in row-major order.
+    expanded_mask = bool_mask
+    for _ in range(self_rank - mask_rank):
+        expanded_mask = op.Unsqueeze(expanded_mask, neg_1)
+    expanded_mask = op.Expand(expanded_mask, op.Shape(self))
+    selected_indices = op.Transpose(op.NonZero(expanded_mask), perm=[1, 0])
+
+    # Broadcast ``values`` to the selection shape ``[num_true, *self.shape[mask_rank:]]``
+    # and flatten it to one update per selected element. This keeps scalar and
+    # broadcastable ``values`` working, matching ``self[mask] = values`` semantics.
+    num_true = op.ReduceSum(op.Cast(op.Reshape(bool_mask, neg_1), to=INT64.dtype), keepdims=1)
+    trailing_shape = op.Slice(op.Shape(self), starts=[mask_rank], ends=[self_rank], axes=[0])
+    selection_shape = op.Concat(num_true, trailing_shape, axis=0)
+    flat_values = op.Reshape(op.Expand(values, selection_shape), neg_1)
+
+    return op.ScatterND(
+        self, selected_indices, flat_values, reduction="add" if accumulate else None
+    )
 
 
 def aten_index_reduce(

{onnxscript-0.7.2.dev20260714 → onnxscript-0.7.2.dev20260731}/onnxscript/rewriter/__init__.py

@@ -43,6 +43,7 @@ from onnxscript.rewriter.rules.common import (
     _fuse_batchnorm,
     _fuse_pad_into_conv,
     _fuse_relus_clips,
+    _materialize_reshape_shape,
     _min_max_to_clip,
     _no_op,
     _redundant_scatter_nd,
@@ -57,6 +58,7 @@ _DEFAULT_REWRITE_RULES: tuple[pattern.RewriteRule, ...] = (
     *_broadcast_to_matmul.rules,
     *_cast_constant_of_shape.rules,
     *_collapse_slices.rules,
+    *_materialize_reshape_shape.rules,
     *_min_max_to_clip.rules,
     *_fuse_relus_clips.rules,
     *_basic_rules.basic_optimization_rules(),

{onnxscript-0.7.2.dev20260714 → onnxscript-0.7.2.dev20260731}/onnxscript/rewriter/rules/common/__init__.py

@@ -26,6 +26,7 @@ __all__ = [
     "max_min_rule",
     "gemm_to_matmul_add_rule",
     "matmul_add_to_gemm_rule",
+    "materialize_reshape_shape_rule",
     "mul_by_1_rule",
     "no_op_cast_rule",
     "no_op_dynamic_scatter_nd_rule",
@@ -102,6 +103,9 @@ from onnxscript.rewriter.rules.common._fuse_relus_clips import (
     successive_relu_rule,
 )
 from onnxscript.rewriter.rules.common._gemm_to_matmul_add import gemm_to_matmul_add_rule
+from onnxscript.rewriter.rules.common._materialize_reshape_shape import (
+    materialize_reshape_shape_rule,
+)
 from onnxscript.rewriter.rules.common._matmul_add_to_gemm import (
     matmul_add_to_gemm_rule,
     transpose_a_matmul_add_to_gemm_rule,

{onnxscript-0.7.2.dev20260714 → onnxscript-0.7.2.dev20260731}/onnxscript/rewriter/rules/common/_collapse_slices.py

@@ -82,7 +82,12 @@ def _same_shape(op, data: ir.Value, slice_output: ir.Value, steps: ir.Value, **_
     if data.shape is None or slice_output.shape is None:
         return False
 
-    if not _ir_utils.is_singleton_value(steps, 1):
+    # All steps must be 1
+    steps_np = _ir_utils.get_numpy_value(steps)
+    if steps_np is not None:
+        if not all(s == 1 for s in steps_np.flat):
+            return False
+    elif not _ir_utils.is_singleton_value(steps, 1):
         return False
 
     return _ir_utils.same_shape(data.shape, slice_output.shape)

onnxscript-0.7.2.dev20260731/onnxscript/rewriter/rules/common/_materialize_reshape_shape.py

@@ -0,0 +1,63 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Materialize Reshape shape input from known output shape.
+
+When symbolic shape inference has been run, a Reshape node may have a known
+output shape even though its shape input is computed dynamically (e.g., via a
+Shape → Cast → Split → Concat chain).  This rule replaces the shape input
+with a concrete constant, allowing the dynamic chain to become dead code and
+be removed by unused-node elimination.
+
+- Fully static output shape → constant with exact dims.
+- Exactly one symbolic dim → replace it with ``-1`` (Reshape infers it).
+"""
+
+from __future__ import annotations
+
+from onnxscript import ir
+from onnxscript.rewriter import _ir_utils as ir_utils
+from onnxscript.rewriter._basics import MatchResult
+from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase, RewriteRuleSet
+
+
+class MaterializeReshapeShape(RewriteRuleClassBase):
+    """Replace a dynamic Reshape shape input with a constant when output shape is known."""
+
+    def pattern(self, op, data, shape):
+        return op.Reshape(data, shape)
+
+    def check(self, context, data: ir.Value, shape: ir.Value) -> MatchResult:
+        check_result = MatchResult()
+
+        # Shape input must not already be a constant
+        if ir_utils.get_numpy_value(shape) is not None:
+            return check_result.fail("Shape input is already a constant.")
+
+        output = context.output_values[0]
+        if output.shape is None:
+            return check_result.fail("Output shape is not known.")
+
+        dims = list(output.shape)
+        sym_count = sum(1 for d in dims if not isinstance(d, int))
+
+        if sym_count <= 1:
+            self._new_dims = [-1 if not isinstance(d, int) else int(d) for d in dims]
+        else:
+            return check_result.fail(
+                f"Output shape has {sym_count} symbolic dims, cannot materialize."
+            )
+
+        return check_result
+
+    def rewrite(self, op, data: ir.Value, shape: ir.Value):
+        new_shape = op.Constant(
+            value=ir.tensor(self._new_dims, dtype=ir.DataType.INT64),
+        )
+        # Use allowzero=1 because the materialized dims are concrete values
+        # from shape inference, so a 0 dim means "actually zero" not "copy from input"
+        return op.Reshape(data, new_shape, allowzero=1)
+
+
+materialize_reshape_shape_rule = MaterializeReshapeShape.rule()
+
+rules = RewriteRuleSet([materialize_reshape_shape_rule])

{onnxscript-0.7.2.dev20260714 → onnxscript-0.7.2.dev20260731/onnxscript.egg-info}/PKG-INFO

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

{onnxscript-0.7.2.dev20260714 → onnxscript-0.7.2.dev20260731}/onnxscript.egg-info/SOURCES.txt

@@ -170,6 +170,7 @@ onnxscript/rewriter/rules/common/_fuse_hardswish.py
 onnxscript/rewriter/rules/common/_fuse_pad_into_conv.py
 onnxscript/rewriter/rules/common/_fuse_relus_clips.py
 onnxscript/rewriter/rules/common/_gemm_to_matmul_add.py
+onnxscript/rewriter/rules/common/_materialize_reshape_shape.py
 onnxscript/rewriter/rules/common/_matmul_add_to_gemm.py
 onnxscript/rewriter/rules/common/_min_max_to_clip.py
 onnxscript/rewriter/rules/common/_no_op.py