onnx-ir 0.0.1__py3-none-any.whl

onnx_ir/_thirdparty/asciichartpy.py

@@ -0,0 +1,313 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+# Copyright © 2016 Igor Kroitor
+#
+# MIT License
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Module to generate ascii charts.
+
+This module provides a single function `plot` that can be used to generate an
+ascii chart from a series of numbers. The chart can be configured via several
+options to tune the output.
+"""
+
+from __future__ import annotations
+
+from math import ceil, floor, isnan
+from typing import Mapping
+
+black = "\033[30m"
+red = "\033[31m"
+green = "\033[32m"
+yellow = "\033[33m"
+blue = "\033[34m"
+magenta = "\033[35m"
+cyan = "\033[36m"
+lightgray = "\033[37m"
+default = "\033[39m"
+darkgray = "\033[90m"
+lightred = "\033[91m"
+lightgreen = "\033[92m"
+lightyellow = "\033[93m"
+lightblue = "\033[94m"
+lightmagenta = "\033[95m"
+lightcyan = "\033[96m"
+white = "\033[97m"
+reset = "\033[0m"
+
+
+__all__ = [
+    "plot",
+    "black",
+    "red",
+    "green",
+    "yellow",
+    "blue",
+    "magenta",
+    "cyan",
+    "lightgray",
+    "default",
+    "darkgray",
+    "lightred",
+    "lightgreen",
+    "lightyellow",
+    "lightblue",
+    "lightmagenta",
+    "lightcyan",
+    "white",
+    "reset",
+]
+
+
+# Python 3.2 has math.isfinite, which could have been used, but to support older
+# versions, this little helper is shorter than having to keep doing not isnan(),
+# plus the double-negative of "not is not a number" is confusing, so this should
+# help with readability.
+def _isnum(n):
+    return not isnan(n)
+
+
+def colored(char, color):
+    if not color:
+        return char
+    else:
+        return color + char + reset
+
+
+_DEFAULT_SYMBOLS = ("┼", "┤", "╶", "╴", "─", "╰", "╭", "╮", "╯", "│")
+
+
+def plot(series, *, bin_edges=None, cfg=None):
+    """Generate an ascii chart for a series of numbers.
+
+    `series` should be a list of ints or floats. Missing data values in the
+    series can be specified as a NaN. In Python versions less than 3.5, use
+    float("nan") to specify an NaN. With 3.5 onwards, use math.nan to specify a
+    NaN.
+
+        >>> series = [1,2,3,4,float("nan"),4,3,2,1]
+        >>> print(plot(series))
+            4.00  ┤  ╭╴╶╮
+            3.00  ┤ ╭╯  ╰╮
+            2.00  ┤╭╯    ╰╮
+            1.00  ┼╯      ╰
+
+    `series` can also be a list of lists to support multiple data series.
+
+        >>> series = [[10,20,30,40,30,20,10], [40,30,20,10,20,30,40]]
+        >>> print(plot(series, cfg={'height': 3}))
+           40.00  ┤╮ ╭╮ ╭
+           30.00  ┤╰╮╯╰╭╯
+           20.00  ┤╭╰╮╭╯╮
+           10.00  ┼╯ ╰╯ ╰
+
+    `bin_edges` is an optional list of bin edges to display on the x-axis. If
+    provided, the x-axis will be labeled with the bin edges. If there are too
+    many bin edges to fit on the x-axis, some labels will be dropped and they
+    will be spaced out evenly to fit the width of the chart.
+    The labels will be formatted using the `x_format` option in `cfg`.
+
+    `cfg` is an optional dictionary of various parameters to tune the appearance
+    of the chart. `min` and `max` will clamp the y-axis and all values:
+
+        >>> series = [1,2,3,4,float("nan"),4,3,2,1]
+        >>> print(plot(series, cfg={'min': 0}))
+            4.00  ┼  ╭╴╶╮
+            3.00  ┤ ╭╯  ╰╮
+            2.00  ┤╭╯    ╰╮
+            1.00  ┼╯      ╰
+            0.00  ┤
+
+        >>> print(plot(series, cfg={'min': 2}))
+            4.00  ┤  ╭╴╶╮
+            3.00  ┤ ╭╯  ╰╮
+            2.00  ┼─╯    ╰─
+
+        >>> print(plot(series, cfg={'min': 2, 'max': 3}))
+            3.00  ┤ ╭─╴╶─╮
+            2.00  ┼─╯    ╰─
+
+    `height` specifies the number of rows the graph should occupy. It can be
+    used to scale down a graph with large data values:
+
+        >>> series = [10,20,30,40,50,40,30,20,10]
+        >>> print(plot(series, cfg={'height': 4}))
+           50.00  ┤   ╭╮
+           40.00  ┤  ╭╯╰╮
+           30.00  ┤ ╭╯  ╰╮
+           20.00  ┤╭╯    ╰╮
+           10.00  ┼╯      ╰
+
+    `format` specifies a Python format string used to format the labels on the
+    y-axis. The default value is "{:8.2f} ". This can be used to remove the
+    decimal point:
+
+        >>> series = [10,20,30,40,50,40,30,20,10]
+        >>> print(plot(series, cfg={'height': 4, 'format':'{:8.0f}'}))
+              50 ┤   ╭╮
+              40 ┤  ╭╯╰╮
+              30 ┤ ╭╯  ╰╮
+              20 ┤╭╯    ╰╮
+              10 ┼╯      ╰
+    """
+    if len(series) == 0:
+        return ""
+
+    if not isinstance(series[0], list):
+        if all(isnan(n) for n in series):
+            return ""
+        else:
+            series = [series]
+
+    if cfg is not None and not isinstance(cfg, Mapping):
+        raise TypeError("cfg must be a dictionary or None")
+
+    cfg = cfg or {}
+
+    colors = cfg.get("colors", [None])
+
+    minimum = cfg.get("min", min(filter(_isnum, [j for i in series for j in i])))
+    maximum = cfg.get("max", max(filter(_isnum, [j for i in series for j in i])))
+
+    symbols = cfg.get("symbols", _DEFAULT_SYMBOLS)
+
+    if minimum > maximum:
+        raise ValueError("The min value cannot exceed the max value.")
+
+    interval = maximum - minimum
+    offset = cfg.get("offset", 3)
+    height = cfg.get("height", interval)
+    ratio = height / interval if interval > 0 else 1
+
+    min2 = floor(minimum * ratio)
+    max2 = ceil(maximum * ratio)
+
+    def clamp(n):
+        return min(max(n, minimum), maximum)
+
+    def scaled(y):
+        return int(round(clamp(y) * ratio) - min2)
+
+    rows = max2 - min2
+
+    width = 0
+    for series_i in series:
+        width = max(width, len(series_i))
+    width += offset
+
+    placeholder = cfg.get("format", "{:8.2f} ")
+    x_placeholder = cfg.get("x_format", "{:4.4f}")
+
+    result = [[" "] * width for i in range(rows + 1)]
+
+    # axis and labels
+    for y in range(min2, max2 + 1):
+        label = placeholder.format(maximum - ((y - min2) * interval / (rows if rows else 1)))
+        result[y - min2][max(offset - len(label), 0)] = label
+        result[y - min2][offset - 1] = symbols[0] if y == 0 else symbols[1]  # zero tick mark
+
+    # first value is a tick mark across the y-axis
+    d0 = series[0][0]
+    if _isnum(d0):
+        result[rows - scaled(d0)][offset - 1] = symbols[0]
+
+    for i, series_i in enumerate(series):
+        color = colors[i % len(colors)]
+
+        # plot the line
+        for x in range(len(series_i) - 1):
+            d0 = series_i[x + 0]
+            d1 = series_i[x + 1]
+
+            if isnan(d0) and isnan(d1):
+                continue
+
+            if isnan(d0) and _isnum(d1):
+                result[rows - scaled(d1)][x + offset] = colored(symbols[2], color)
+                continue
+
+            if _isnum(d0) and isnan(d1):
+                result[rows - scaled(d0)][x + offset] = colored(symbols[3], color)
+                continue
+
+            y0 = scaled(d0)
+            y1 = scaled(d1)
+            if y0 == y1:
+                result[rows - y0][x + offset] = colored(symbols[4], color)
+                continue
+
+            result[rows - y1][x + offset] = (
+                colored(symbols[5], color) if y0 > y1 else colored(symbols[6], color)
+            )
+            result[rows - y0][x + offset] = (
+                colored(symbols[7], color) if y0 > y1 else colored(symbols[8], color)
+            )
+
+            start = min(y0, y1) + 1
+            end = max(y0, y1)
+            for y in range(start, end):
+                result[rows - y][x + offset] = colored(symbols[9], color)
+
+    the_plot = "\n".join(["".join(row).rstrip() for row in result])
+
+    if bin_edges is None or len(bin_edges) == 0:
+        return the_plot
+
+    # Plot x axis labels
+    current_location = 0
+    # Compute the amount of leading space for the first x-label using the old label size
+    leading_space = offset + len(label)
+    # Obtain the first x-label to compute its size
+    x_label = x_placeholder.format(bin_edges[0])
+    # Initialize the x-label text with the leading space. We allow the first label to
+    # recess so that the center of it is aligned with the first tick mark.
+    x_label_size = len(x_label)
+    x_leading_space = max(0, leading_space - x_label_size)
+
+    x_labels = []
+    # This is the amount of space we have to fit the x-labels. It can overflow the width
+    # by half of the x-label size
+    workable_width = width + x_label_size // 2
+    # Compute the spacing between x-labels
+    # If we fit labels and space them by 2 characters, we can fit this many labels:
+    min_spacing = 2
+    num_labels_can_fit = width // (x_label_size + min_spacing)
+    labels_count = len(bin_edges)
+    # Find out the actual number of labels we need to display
+    num_labels_to_display = min(labels_count, num_labels_can_fit)
+    num_spaces = num_labels_to_display - 1
+    spacing = max(
+        min_spacing,
+        (workable_width - num_labels_to_display * x_label_size) // num_spaces,
+    )
+    # Now start placing labels
+    while current_location < workable_width:
+        # Find the current label that would be suitable for the current location
+        bin_index = int((current_location / workable_width) * labels_count)
+        x_label = x_placeholder.format(bin_edges[bin_index])
+        x_labels.append(x_label)
+        # Move to the next location
+        current_location += len(x_label) + spacing
+    # Create the x-label row
+    x_labels_text = " " * x_leading_space + (" " * spacing).join(x_labels)
+
+    return the_plot + "\n" + x_labels_text

onnx_ir/_type_casting.py

@@ -0,0 +1,91 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Numpy utilities for non-native type operation."""
+# TODO(justinchuby): Upstream the logic to onnx
+
+from __future__ import annotations
+
+import typing
+from typing import Sequence
+
+import ml_dtypes
+import numpy as np
+
+if typing.TYPE_CHECKING:
+    import numpy.typing as npt
+
+
+def pack_int4(array: np.ndarray) -> npt.NDArray[np.uint8]:
+    """Convert a numpy array to flatten, packed int4/uint4. Elements must be in the correct range."""
+    # Create a 1D copy
+    array_flat = array.ravel().view(np.uint8).copy()
+    size = array.size
+    odd_sized = size % 2 == 1
+    if odd_sized:
+        array_flat.resize([size + 1], refcheck=False)
+    array_flat &= 0x0F
+    array_flat[1::2] <<= 4
+    return array_flat[0::2] | array_flat[1::2]  # type: ignore[return-type]
+
+
+def _unpack_uint4_as_uint8(
+    data: npt.NDArray[np.uint8], dims: Sequence[int]
+) -> npt.NDArray[np.uint8]:
+    """Convert a packed uint4 array to unpacked uint4 array represented as uint8.
+
+    Args:
+        data: A numpy array.
+        dims: The dimensions are used to reshape the unpacked buffer.
+
+    Returns:
+        A numpy array of int8/uint8 reshaped to dims.
+    """
+    result = np.empty([data.size * 2], dtype=data.dtype)
+    array_low = data & np.uint8(0x0F)
+    array_high = data & np.uint8(0xF0)
+    array_high >>= np.uint8(4)
+    result[0::2] = array_low
+    result[1::2] = array_high
+    if result.size == np.prod(dims) + 1:
+        # handle single-element padding due to odd number of elements
+        result = result[:-1]
+    result.resize(dims, refcheck=False)
+    return result
+
+
+def unpack_uint4(
+    data: npt.NDArray[np.uint8], dims: Sequence[int]
+) -> npt.NDArray[ml_dtypes.uint4]:
+    """Convert a packed uint4 array to unpacked uint4 array represented as uint8.
+
+    Args:
+        data: A numpy array.
+        dims: The dimensions are used to reshape the unpacked buffer.
+
+    Returns:
+        A numpy array of int8/uint8 reshaped to dims.
+    """
+    return _unpack_uint4_as_uint8(data, dims).view(ml_dtypes.uint4)
+
+
+def _extend_int4_sign_bits(x: npt.NDArray[np.uint8]) -> npt.NDArray[np.int8]:
+    """Extend 4-bit signed integer to 8-bit signed integer."""
+    return np.where((x >> 3) == 0, x, x | 0xF0).astype(np.int8)
+
+
+def unpack_int4(
+    data: npt.NDArray[np.uint8], dims: Sequence[int]
+) -> npt.NDArray[ml_dtypes.int4]:
+    """Convert a packed (signed) int4 array to unpacked int4 array represented as int8.
+
+    The sign bit is extended to the most significant bit of the int8.
+
+    Args:
+        data: A numpy array.
+        dims: The dimensions are used to reshape the unpacked buffer.
+
+    Returns:
+        A numpy array of int8 reshaped to dims.
+    """
+    unpacked = _unpack_uint4_as_uint8(data, dims)
+    return _extend_int4_sign_bits(unpacked).view(ml_dtypes.int4)

onnx_ir/convenience.py

@@ -0,0 +1,32 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Convenience methods for constructing and manipulating the IR."""
+
+from __future__ import annotations
+
+__all__ = [
+    "convert_attribute",
+    "convert_attributes",
+    "replace_all_uses_with",
+    "replace_nodes_and_values",
+]
+
+from onnx_ir._convenience import (
+    convert_attribute,
+    convert_attributes,
+    replace_all_uses_with,
+    replace_nodes_and_values,
+)
+
+# NOTE: Do not implement any other functions in this module.
+# implement them in the _convenience module and import them here instead.
+
+
+def __set_module() -> None:
+    """Set the module of all functions in this module to this public module."""
+    global_dict = globals()
+    for name in __all__:
+        global_dict[name].__module__ = __name__
+
+
+__set_module()

onnx_ir/passes/__init__.py

@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+__all__ = [
+    "PassBase",
+    "PassResult",
+    "PassManager",
+    # Errors
+    "InvariantError",
+    "PreconditionError",
+    "PostconditionError",
+    "PassError",
+]
+
+from onnx_ir.passes._pass_infra import (
+    InvariantError,
+    PassBase,
+    PassError,
+    PassManager,
+    PassResult,
+    PostconditionError,
+    PreconditionError,
+)
+
+
+def __set_module() -> None:
+    """Set the module of all functions in this module to this public module."""
+    global_dict = globals()
+    for name in __all__:
+        global_dict[name].__module__ = __name__
+
+
+__set_module()

onnx_ir/passes/_pass_infra.py

@@ -0,0 +1,172 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+# This module implements some APIs described in
+# https://pytorch.org/executorch/stable/compiler-custom-compiler-passes.html
+# for the ONNX IR.
+# The classes {PassResult and PassManager} are derived from
+# https://github.com/pytorch/pytorch/blob/1e47c7b11b312b47a621efd547f5c90081f0d9cb/torch/fx/passes/infra/pass_base.py#L12
+# and
+# https://github.com/pytorch/pytorch/blob/1e47c7b11b312b47a621efd547f5c90081f0d9cb/torch/fx/passes/infra/pass_manager.py#L147
+# The original code is licensed under the PyTorch License https://github.com/pytorch/pytorch/blob/main/LICENSE
+
+"""Passes infrastructure for the IR."""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+from typing import Sequence
+
+__all__ = [
+    "PassBase",
+    "PassManager",
+    "PassResult",
+    # Errors
+    "InvariantError",
+    "PreconditionError",
+    "PostconditionError",
+    "PassError",
+]
+
+import abc
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+class InvariantError(Exception):
+    """Raised when an invariant is violated."""
+
+
+class PreconditionError(InvariantError):
+    """Raised when a precondition is violated."""
+
+
+class PostconditionError(InvariantError):
+    """Raised when a postcondition is violated."""
+
+
+class PassError(RuntimeError):
+    """Raised when an error occurs during a pass."""
+
+
+@dataclasses.dataclass
+class PassResult:
+    """Result of a pass.
+
+    Attributes:
+        model: The transformed model.
+        modified: Whether the model was modified.
+    """
+
+    model: ir.Model
+    modified: bool
+
+
+class PassBase(abc.ABC):
+    """Base class for all passes.
+
+    Class attributes:
+        in_place: Whether the pass modifies the model in place.
+    """
+
+    in_place: bool = True
+
+    def __call__(self, model: ir.Model) -> PassResult:
+        return self.call(model)
+
+    @abc.abstractmethod
+    def call(self, model: ir.Model) -> PassResult:
+        """The main entry point for the pass."""
+        ...
+
+    def requires(self, model: ir.Model) -> None:
+        """Pre-conditions for the pass.
+
+        This is optional to implement, will be called before call() if run by a pass manager.
+        """
+        del model  # Unused
+
+    def ensures(self, model: ir.Model) -> None:
+        """Post-conditions for the pass.
+
+        This is optional to implement, will be called after call() if run by a pass manager.
+        """
+        del model  # Unused
+
+
+class PassManager:
+    """Pass manager for the IR.
+
+    The PassManager is a callable that runs a sequence of passes on a model.
+
+    Attributes:
+        passes: The passes to run.
+        check_invariants: Whether to check invariants before and after each pass.
+        steps: The number of times to run the passes.
+    """
+
+    def __init__(
+        self,
+        passes: Sequence[PassBase],
+        check_invariants: bool = False,
+        steps: int = 1,
+    ):
+        # TODO(justinchuby): Implement constraints
+        self.passes = list(passes)
+        self.check_invariants = check_invariants
+        self.steps = steps
+
+    def __call__(self, model: ir.Model) -> PassResult:
+        """Run the set of passes `steps` number of times or until the graph stops changing."""
+        overall_modified = False
+        for step in range(self.steps):
+            step_result = self._run_one_step(model, step)
+            model = step_result.model
+            modified = step_result.modified
+            overall_modified = overall_modified or modified
+            # If the graph no longer changes, then we can stop running these passes
+            if not modified:
+                logger.info("PassManager: No more graph changes detected after step %s", step)
+                break
+        return PassResult(model, overall_modified)
+
+    def _run_one_step(self, model: ir.Model, step: int) -> PassResult:
+        modified = False
+        for i, pass_ in enumerate(self.passes):
+            logger.debug("Running the %s-th pass '%s', (step %s)", i, pass_, step)
+
+            # 1. Check preconditions
+            if self.check_invariants:
+                try:
+                    pass_.requires(model)
+                except Exception as e:
+                    raise PreconditionError(f"Pre-condition failed for {pass_}") from e
+
+            # 2. Run the pass
+            try:
+                pass_result = pass_(model)
+            except Exception as e:
+                prev_pass_names = [str(p) for p in self.passes[:i]]
+                raise PassError(
+                    f"An error occurred when running the '{pass_}' pass after the "
+                    f"following passes: {prev_pass_names} during step {step}"
+                ) from e
+            if not isinstance(pass_result, PassResult):
+                raise TypeError(
+                    f"The result of the pass {pass_} should be type PassResult."
+                    "Please create one with ir.passes.PassResult()."
+                )
+
+            model = pass_result.model
+            modified = modified or pass_result.modified
+
+            # 3. Check postconditions
+            if self.check_invariants:
+                try:
+                    pass_.ensures(model)
+                except Exception as e:
+                    raise PostconditionError(f"Post-condition failed for {pass_}") from e
+        return PassResult(model, modified)