onnx2fx 0.0.0__py3-none-any.whl

onnx2fx/ops/reduction.py

@@ -0,0 +1,534 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Reduction operators."""
+
+from typing import TYPE_CHECKING
+
+import onnx
+import torch
+
+from ..op_registry import register
+from ..utils.attributes import get_attribute
+from ..utils.op_helpers import get_attribute_or_input
+
+if TYPE_CHECKING:
+    from ..graph_builder import GraphBuilder
+
+
+def _get_reduction_axes(
+    node: onnx.NodeProto, builder: "GraphBuilder"
+) -> list[int] | torch.fx.Node | None:
+    """Get axes for reduction, handling both attribute and input formats.
+
+    In opset < 13, axes is an attribute.
+    In opset 13-17, axes can be an attribute or an optional input.
+    In opset 18+, axes is an optional input only.
+    """
+    return get_attribute_or_input(
+        builder,
+        node,
+        attr_name="axes",
+        input_index=1,
+        opset_version=builder.opset_version,
+        attr_allowed_until=17,
+        input_allowed_since=13,
+        default=None,
+    )
+
+
+@register("ReduceSum")
+def reduce_sum(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Sum reduction."""
+    x = builder.get_value(node.input[0])
+    axes = _get_reduction_axes(node, builder)
+    keepdims = get_attribute(node, "keepdims", 1)
+    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)
+
+    def _reduce_sum(t, axes, keepdims, noop_with_empty_axes):
+        # Handle empty axes list
+        if isinstance(axes, (list, tuple)) and len(axes) == 0:
+            if noop_with_empty_axes:
+                return t
+            # Empty axes with noop=False means reduce all dimensions
+            axes = None
+        if isinstance(axes, torch.Tensor) and axes.numel() == 0:
+            if noop_with_empty_axes:
+                return t
+            axes = None
+        if axes is None:
+            result = torch.sum(t)
+            if keepdims:
+                # Reshape to have all dimensions as 1
+                result = result.reshape([1] * t.ndim)
+            return result
+        if isinstance(axes, torch.Tensor):
+            axes = tuple(axes.tolist())
+        elif isinstance(axes, (list, tuple)):
+            axes = tuple(axes)
+        return torch.sum(t, dim=axes, keepdim=keepdims)
+
+    return builder.call_function(
+        _reduce_sum, args=(x, axes, bool(keepdims), bool(noop_with_empty_axes))
+    )
+
+
+@register("ReduceMean")
+def reduce_mean(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Mean reduction."""
+    x = builder.get_value(node.input[0])
+    axes = _get_reduction_axes(node, builder)
+    keepdims = get_attribute(node, "keepdims", 1)
+    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)
+
+    def _reduce_mean(t, axes, keepdims, noop_with_empty_axes):
+        # Handle empty axes list
+        if isinstance(axes, (list, tuple)) and len(axes) == 0:
+            if noop_with_empty_axes:
+                return t
+            axes = None
+        if isinstance(axes, torch.Tensor) and axes.numel() == 0:
+            if noop_with_empty_axes:
+                return t
+            axes = None
+        if axes is None:
+            result = torch.mean(t)
+            if keepdims:
+                result = result.reshape([1] * t.ndim)
+            return result
+        if isinstance(axes, torch.Tensor):
+            axes = tuple(axes.tolist())
+        elif isinstance(axes, (list, tuple)):
+            axes = tuple(axes)
+        return torch.mean(t, dim=axes, keepdim=keepdims)
+
+    return builder.call_function(
+        _reduce_mean, args=(x, axes, bool(keepdims), bool(noop_with_empty_axes))
+    )
+
+
+@register("ReduceMax")
+def reduce_max(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Max reduction."""
+    x = builder.get_value(node.input[0])
+    axes = _get_reduction_axes(node, builder)
+    keepdims = get_attribute(node, "keepdims", 1)
+    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)
+
+    def _reduce_max(t, axes, keepdims, noop_with_empty_axes):
+        # Handle empty axes list
+        if isinstance(axes, (list, tuple)) and len(axes) == 0:
+            if noop_with_empty_axes:
+                return t
+            axes = None
+        if isinstance(axes, torch.Tensor) and axes.numel() == 0:
+            if noop_with_empty_axes:
+                return t
+            axes = None
+
+        if axes is None:
+            # Reduce over all dimensions
+            if t.numel() == 0:
+                # Empty tensor: return -inf with proper shape
+                if keepdims:
+                    return torch.full(
+                        [1] * t.ndim, float("-inf"), dtype=t.dtype, device=t.device
+                    )
+                return torch.tensor(float("-inf"), dtype=t.dtype, device=t.device)
+            result = t.max()
+            if keepdims:
+                result = result.reshape([1] * t.ndim)
+            return result
+
+        if isinstance(axes, torch.Tensor):
+            axes = axes.tolist()
+        if isinstance(axes, list) and len(axes) == 1:
+            axes = axes[0]
+        if isinstance(axes, int):
+            # Check for empty dimension
+            if t.shape[axes] == 0:
+                new_shape = list(t.shape)
+                if keepdims:
+                    new_shape[axes] = 1
+                else:
+                    new_shape.pop(axes)
+                return torch.full(
+                    new_shape, float("-inf"), dtype=t.dtype, device=t.device
+                )
+            return t.max(dim=axes, keepdim=keepdims).values
+        # Multiple axes: reduce sequentially
+        result = t
+        for axis in sorted(axes, reverse=True):
+            if result.shape[axis] == 0:
+                new_shape = list(result.shape)
+                if keepdims:
+                    new_shape[axis] = 1
+                else:
+                    new_shape.pop(axis)
+                result = torch.full(
+                    new_shape, float("-inf"), dtype=result.dtype, device=result.device
+                )
+            else:
+                result = result.max(dim=axis, keepdim=keepdims).values
+        return result
+
+    return builder.call_function(
+        _reduce_max, args=(x, axes, bool(keepdims), bool(noop_with_empty_axes))
+    )
+
+
+@register("ReduceMin")
+def reduce_min(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Min reduction."""
+    x = builder.get_value(node.input[0])
+    axes = _get_reduction_axes(node, builder)
+    keepdims = get_attribute(node, "keepdims", 1)
+    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)
+
+    def _reduce_min(t, axes, keepdims, noop_with_empty_axes):
+        # Handle empty axes list
+        if isinstance(axes, (list, tuple)) and len(axes) == 0:
+            if noop_with_empty_axes:
+                return t
+            axes = None
+        if isinstance(axes, torch.Tensor) and axes.numel() == 0:
+            if noop_with_empty_axes:
+                return t
+            axes = None
+
+        if axes is None:
+            # Reduce over all dimensions
+            if t.numel() == 0:
+                # Empty tensor: return inf with proper shape
+                if keepdims:
+                    return torch.full(
+                        [1] * t.ndim, float("inf"), dtype=t.dtype, device=t.device
+                    )
+                return torch.tensor(float("inf"), dtype=t.dtype, device=t.device)
+            result = t.min()
+            if keepdims:
+                result = result.reshape([1] * t.ndim)
+            return result
+
+        if isinstance(axes, torch.Tensor):
+            axes = axes.tolist()
+        if isinstance(axes, list) and len(axes) == 1:
+            axes = axes[0]
+        if isinstance(axes, int):
+            # Check for empty dimension
+            if t.shape[axes] == 0:
+                new_shape = list(t.shape)
+                if keepdims:
+                    new_shape[axes] = 1
+                else:
+                    new_shape.pop(axes)
+                return torch.full(
+                    new_shape, float("inf"), dtype=t.dtype, device=t.device
+                )
+            return t.min(dim=axes, keepdim=keepdims).values
+        # Multiple axes: reduce sequentially
+        result = t
+        for axis in sorted(axes, reverse=True):
+            if result.shape[axis] == 0:
+                new_shape = list(result.shape)
+                if keepdims:
+                    new_shape[axis] = 1
+                else:
+                    new_shape.pop(axis)
+                result = torch.full(
+                    new_shape, float("inf"), dtype=result.dtype, device=result.device
+                )
+            else:
+                result = result.min(dim=axis, keepdim=keepdims).values
+        return result
+
+    return builder.call_function(
+        _reduce_min, args=(x, axes, bool(keepdims), bool(noop_with_empty_axes))
+    )
+
+
+@register("ReduceProd")
+def reduce_prod(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Product reduction."""
+    x = builder.get_value(node.input[0])
+    axes = _get_reduction_axes(node, builder)
+    keepdims = get_attribute(node, "keepdims", 1)
+    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)
+
+    def _reduce_prod(t, axes, keepdims, noop_with_empty_axes):
+        # Handle empty axes list
+        if isinstance(axes, (list, tuple)) and len(axes) == 0:
+            if noop_with_empty_axes:
+                return t
+            axes = None
+        if isinstance(axes, torch.Tensor) and axes.numel() == 0:
+            if noop_with_empty_axes:
+                return t
+            axes = None
+
+        if axes is None:
+            result = torch.prod(t)
+            if keepdims:
+                result = result.reshape([1] * t.ndim)
+            return result
+
+        if isinstance(axes, torch.Tensor):
+            axes = axes.tolist()
+        if isinstance(axes, list) and len(axes) == 1:
+            axes = axes[0]
+        if isinstance(axes, int):
+            return torch.prod(t, dim=axes, keepdim=keepdims)
+        # Multiple axes: reduce sequentially
+        result = t
+        for axis in sorted(axes, reverse=True):
+            result = torch.prod(result, dim=axis, keepdim=keepdims)
+        return result
+
+    return builder.call_function(
+        _reduce_prod, args=(x, axes, bool(keepdims), bool(noop_with_empty_axes))
+    )
+
+
+@register("ReduceL1")
+def reduce_l1(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """L1 norm reduction."""
+    x = builder.get_value(node.input[0])
+    axes = _get_reduction_axes(node, builder)
+    keepdims = get_attribute(node, "keepdims", 1)
+
+    def _reduce_l1(t, axes, keepdims):
+        abs_t = torch.abs(t)
+        if axes is None:
+            return torch.sum(abs_t)
+        if isinstance(axes, torch.Tensor):
+            axes = tuple(axes.tolist())
+        elif isinstance(axes, (list, tuple)):
+            axes = tuple(axes)
+        return torch.sum(abs_t, dim=axes, keepdim=keepdims)
+
+    return builder.call_function(_reduce_l1, args=(x, axes, bool(keepdims)))
+
+
+@register("ReduceL2")
+def reduce_l2(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """L2 norm reduction."""
+    x = builder.get_value(node.input[0])
+    axes = _get_reduction_axes(node, builder)
+    keepdims = get_attribute(node, "keepdims", 1)
+
+    def _reduce_l2(t, axes, keepdims):
+        if axes is None:
+            return torch.norm(t)
+        if isinstance(axes, torch.Tensor):
+            axes = tuple(axes.tolist())
+        elif isinstance(axes, (list, tuple)):
+            axes = tuple(axes)
+        return torch.norm(t, dim=axes, keepdim=keepdims)
+
+    return builder.call_function(_reduce_l2, args=(x, axes, bool(keepdims)))
+
+
+@register("ReduceLogSum")
+def reduce_log_sum(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Log of sum reduction."""
+    x = builder.get_value(node.input[0])
+    axes = _get_reduction_axes(node, builder)
+    keepdims = get_attribute(node, "keepdims", 1)
+
+    def _reduce_log_sum(t, axes, keepdims):
+        if axes is None:
+            return torch.log(torch.sum(t))
+        if isinstance(axes, torch.Tensor):
+            axes = tuple(axes.tolist())
+        elif isinstance(axes, (list, tuple)):
+            axes = tuple(axes)
+        return torch.log(torch.sum(t, dim=axes, keepdim=keepdims))
+
+    return builder.call_function(_reduce_log_sum, args=(x, axes, bool(keepdims)))
+
+
+@register("ReduceLogSumExp")
+def reduce_log_sum_exp(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """LogSumExp reduction."""
+    x = builder.get_value(node.input[0])
+    axes = _get_reduction_axes(node, builder)
+    keepdims = get_attribute(node, "keepdims", 1)
+
+    def _reduce_log_sum_exp(t, axes, keepdims):
+        if axes is None:
+            return torch.logsumexp(t, dim=tuple(range(t.dim())))
+        if isinstance(axes, torch.Tensor):
+            axes = tuple(axes.tolist())
+        elif isinstance(axes, (list, tuple)):
+            axes = tuple(axes)
+        return torch.logsumexp(t, dim=axes, keepdim=keepdims)
+
+    return builder.call_function(_reduce_log_sum_exp, args=(x, axes, bool(keepdims)))
+
+
+@register("ReduceSumSquare")
+def reduce_sum_square(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Sum of squares reduction."""
+    x = builder.get_value(node.input[0])
+    axes = _get_reduction_axes(node, builder)
+    keepdims = get_attribute(node, "keepdims", 1)
+
+    def _reduce_sum_square(t, axes, keepdims):
+        sq = torch.square(t)
+        if axes is None:
+            return torch.sum(sq)
+        if isinstance(axes, torch.Tensor):
+            axes = tuple(axes.tolist())
+        elif isinstance(axes, (list, tuple)):
+            axes = tuple(axes)
+        return torch.sum(sq, dim=axes, keepdim=keepdims)
+
+    return builder.call_function(_reduce_sum_square, args=(x, axes, bool(keepdims)))
+
+
+# =============================================================================
+# ArgMax/ArgMin operators
+# =============================================================================
+
+
+def _make_arg_extremum_handler(torch_fn):
+    """Factory for ArgMax/ArgMin operator handlers."""
+
+    def _arg_extremum(t, axis, keepdims, select_last_index):
+        if select_last_index:
+            flipped = torch.flip(t, [axis])
+            idx = torch_fn(flipped, dim=axis, keepdim=keepdims)
+            return t.size(axis) - 1 - idx
+        return torch_fn(t, dim=axis, keepdim=keepdims)
+
+    def handler(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+        x = builder.get_value(node.input[0])
+        axis = get_attribute(node, "axis", 0)
+        keepdims = get_attribute(node, "keepdims", 1)
+        select_last_index = get_attribute(node, "select_last_index", 0)
+        return builder.call_function(
+            _arg_extremum, args=(x, axis, bool(keepdims), bool(select_last_index))
+        )
+
+    return handler
+
+
+@register("ArgMax")
+def argmax(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Index of maximum value."""
+    return _make_arg_extremum_handler(torch.argmax)(builder, node)
+
+
+@register("ArgMin")
+def argmin(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Index of minimum value."""
+    return _make_arg_extremum_handler(torch.argmin)(builder, node)
+
+
+# =============================================================================
+# Cumulative and TopK operators
+# =============================================================================
+
+
+@register("CumSum")
+def cumsum(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Cumulative sum."""
+    x = builder.get_value(node.input[0])
+    axis = builder.get_value(node.input[1])
+
+    exclusive = get_attribute(node, "exclusive", 0)
+    reverse = get_attribute(node, "reverse", 0)
+
+    def _cumsum(x, axis, exclusive, reverse):
+        ax = axis.item() if isinstance(axis, torch.Tensor) else axis
+
+        if reverse:
+            x = torch.flip(x, [int(ax)])
+
+        result = torch.cumsum(x, dim=int(ax))
+
+        if exclusive:
+            # Shift by one and pad with zero
+            pad_shape = list(x.shape)
+            pad_shape[int(ax)] = 1
+            zero_pad = torch.zeros(pad_shape, dtype=x.dtype, device=x.device)
+            result = torch.cat(
+                [zero_pad, result.narrow(int(ax), 0, x.shape[int(ax)] - 1)], dim=int(ax)
+            )
+
+        if reverse:
+            result = torch.flip(result, [int(ax)])
+
+        return result
+
+    return builder.call_function(_cumsum, args=(x, axis, exclusive, reverse))
+
+
+@register("TopK")
+def topk(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Find top K values and indices."""
+    x = builder.get_value(node.input[0])
+    k = builder.get_value(node.input[1])
+
+    axis = get_attribute(node, "axis", -1)
+    largest = get_attribute(node, "largest", 1)
+    sorted_ = get_attribute(node, "sorted", 1)
+
+    def _topk(x, k, axis, largest, sorted_):
+        k_val = k.item() if isinstance(k, torch.Tensor) else k
+        k_val = int(k_val)
+
+        # Handle unsupported dtypes (e.g., uint64) by converting to int64
+        original_dtype = x.dtype
+        needs_conversion = original_dtype == torch.uint64
+        if needs_conversion:
+            x = x.to(torch.int64)
+
+        # ONNX TopK requires stable sorting: for equal values, the element
+        # with lower index appears first. PyTorch's topk is not stable.
+        # We achieve stability by using argsort on a composite key.
+        # Create indices tensor for tie-breaking
+        size = x.shape[axis]
+        # Create indices [0, 1, 2, ..., size-1] along the specified axis
+        indices_shape = [1] * x.ndim
+        indices_shape[axis] = size
+        idx = torch.arange(size, device=x.device, dtype=x.dtype).view(indices_shape)
+        idx = idx.expand_as(x)
+
+        # Scale values so that the index becomes the tiebreaker
+        # For largest=True: negate values, sort ascending, lower index wins
+        # For largest=False: use values directly, sort ascending, lower index wins
+        if bool(largest):
+            # Negate so that larger values become smaller (for ascending sort)
+            # Add small offset based on index to break ties (lower index = smaller offset)
+            sort_values = -x
+        else:
+            sort_values = x
+
+        # Use argsort with stable=True for stable sorting
+        sorted_indices = torch.argsort(sort_values, dim=axis, stable=True)
+
+        # Take top k indices
+        # Narrow to first k elements along axis
+        top_k_indices = torch.narrow(sorted_indices, axis, 0, k_val)
+
+        # Gather values using the indices
+        values = torch.gather(x, axis, top_k_indices)
+
+        # If sorted=False, the order is undefined, but we still use stable order
+        # The indices should be the original indices
+        indices = top_k_indices
+
+        # Convert values back to original dtype if needed
+        if needs_conversion:
+            values = values.to(original_dtype)
+
+        return values, indices
+
+    result = builder.call_function(_topk, args=(x, k, axis, largest, sorted_))
+
+    # Handle multiple outputs
+    for i, output_name in enumerate(node.output):
+        if output_name:
+            idx_node = builder.call_function(lambda t, idx: t[idx], args=(result, i))
+            builder.env[output_name] = idx_node
+
+    return result