onnx-ir 0.1.15__py3-none-any.whl

onnx_ir/_graph_containers.py

@@ -0,0 +1,373 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Tracked containers for graph."""
+
+# pylint: disable=protected-access
+
+from __future__ import annotations
+
+__all__ = [
+    "GraphInputs",
+    "GraphOutputs",
+]
+
+import collections
+import logging
+from collections.abc import Iterable, Sequence
+from typing import SupportsIndex, TypeVar
+
+import onnx_ir
+from onnx_ir import _core, _protocols
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+class _GraphIO(collections.UserList["_core.Value"]):
+    """The inputs and outputs of a Graph."""
+
+    def __init__(self, graph: _core.Graph, initlist=None):
+        self._graph = graph
+        # Use a ref counter to track the number of references to each value
+        # in the input/output list. This is used to determine when to unset the graph
+        # reference in the value.
+        # Even though a duplicated value is invalid in inputs and not recommended in outputs,
+        # it is still possible to have duplicated inputs/outputs in an ONNX graph so we
+        # need to properly handle this case and maintain the graph reference properly.
+        self._ref_counter: collections.Counter[_core.Value] = collections.Counter()
+        if initlist is not None:
+            initlist = tuple(initlist)  # Create a copy in case initlist is a generator
+            for value in initlist:
+                self._set_graph(value)
+        super().__init__(initlist)
+        self._check_invariance()
+
+    def _check_invariance(self) -> None:
+        """Check the invariance of the graph."""
+        raise NotImplementedError
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        raise NotImplementedError
+
+    def _maybe_unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        raise NotImplementedError
+
+    def append(self, item: _core.Value) -> None:
+        """Add a new input to the graph."""
+        # Perform checks first in _set_graph before modifying the data structure
+        self._set_graph(item)
+        super().append(item)
+        self._check_invariance()
+
+    def extend(self, other) -> None:
+        """Extend the list of inputs or outputs."""
+        other = tuple(other)
+        for item in other:
+            self._set_graph(item)
+        super().extend(other)
+
+    def insert(self, i: int, item: _core.Value) -> None:
+        """Insert an input/output to the graph."""
+        super().insert(i, item)
+        self._set_graph(item)
+        self._check_invariance()
+
+    def pop(self, i: int = -1) -> _core.Value:
+        """Remove an input/output from the graph."""
+        value = super().pop(i)
+        self._maybe_unset_graph(value)
+        self._check_invariance()
+        return value
+
+    def remove(self, item: _core.Value) -> None:
+        """Remove an input/output from the graph."""
+        super().remove(item)
+        self._maybe_unset_graph(item)
+        self._check_invariance()
+
+    def clear(self) -> None:
+        """Clear the list."""
+        for value in self.data:
+            self._maybe_unset_graph(value)
+        super().clear()
+
+    def copy(self) -> list[_core.Value]:
+        """Return a shallow copy of the list."""
+        # This is a shallow copy, so the values are not copied, just the references
+        return self.data.copy()
+
+    def __setitem__(self, i, item) -> None:
+        """Replace an input/output to the node."""
+        if isinstance(item, Iterable) and isinstance(i, slice):
+            # Modify a slice of the list
+            for value in self.data[i]:
+                self._maybe_unset_graph(value)
+            for value in item:
+                self._set_graph(value)
+            super().__setitem__(i, item)
+            self._check_invariance()
+            return
+        elif isinstance(i, SupportsIndex):
+            # Replace a single item
+            self._maybe_unset_graph(self.data[i])
+            self._set_graph(item)
+            super().__setitem__(i, item)
+            self._check_invariance()
+            return
+
+        raise TypeError(f"Invalid types for __setitem__: {type(i)} and {type(item)}")
+
+    def __getitem__(self, i):
+        """Get an input/output from the graph."""
+        return self.data[i]
+
+    def _unimplemented(self, *_args, **_kwargs):
+        """Unimplemented method."""
+        raise RuntimeError("Method is not supported")
+
+    __add__ = _unimplemented
+    __radd__ = _unimplemented
+    __iadd__ = _unimplemented
+    __mul__ = _unimplemented
+    __rmul__ = _unimplemented
+
+
+class GraphInputs(_GraphIO):
+    """The inputs of a Graph."""
+
+    def _check_invariance(self) -> None:
+        """Check the invariance of the graph."""
+        if not onnx_ir.DEBUG:
+            return
+        for value in self.data:
+            if value._graph is self._graph:
+                continue
+            raise ValueError(
+                f"Invariance error: Value '{value}' is not an input of the graph: {self._graph!r}"
+            )
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        if value._graph is not None and value._graph is not self._graph:
+            raise ValueError(
+                f"Value '{value}' is already owned by a different graph. Please remove the value from the previous graph first"
+            )
+        if value.producer() is not None:
+            raise ValueError(
+                f"Value '{value}' is produced by a node and cannot be an input to the graph. Please create new Values for graph inputs"
+            )
+        self._ref_counter[value] += 1
+        value._is_graph_input = True
+        value._graph = self._graph
+
+    def _maybe_unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        assert value._graph is self._graph, "Bug: value does not belong to the graph"
+        self._ref_counter[value] -= 1
+        if self._ref_counter[value] > 0:
+            # The value is still used by another graph input
+            return
+        value._is_graph_input = False
+        if value._owned_by_graph():
+            # Keep the graph reference if the value is still an input or an initializer
+            return
+        value._graph = None
+
+
+class GraphOutputs(_GraphIO):
+    """The outputs of a Graph."""
+
+    def _check_invariance(self) -> None:
+        """Check the invariance of the graph."""
+        if not onnx_ir.DEBUG:
+            return
+        for value in self.data:
+            if value._graph is self._graph:
+                continue
+            raise ValueError(
+                f"Invariance error: Value '{value}' is not an output of the graph: {self._graph!r}"
+            )
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        if value._graph is not None and value._graph is not self._graph:
+            raise ValueError(
+                f"Value '{value}' is already an output of a different graph. Please remove the value from the previous graph first"
+            )
+        self._ref_counter[value] += 1
+        value._is_graph_output = True
+        value._graph = self._graph
+
+    def _maybe_unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        assert value._graph is self._graph, "Bug: value does not belong to the graph"
+        self._ref_counter[value] -= 1
+        if self._ref_counter[value] > 0:
+            # The value is still used by another graph input
+            return
+        value._is_graph_output = False
+        if value._owned_by_graph():
+            # Keep the graph reference if the value is still an input or an initializer
+            return
+        value._graph = None
+
+
+class GraphInitializers(collections.UserDict[str, "_core.Value"]):
+    """The initializers of a Graph as ``dict[str, Value]`` with additional mutation methods."""
+
+    def __init__(self, graph: _core.Graph, dict=None, /, **kwargs):
+        # Perform checks first in _set_graph before modifying the data structure with super().__init__()
+        data = {}
+        if dict is not None:
+            data.update(dict)
+        if kwargs:
+            data.update(kwargs)
+        self._graph = graph
+        for value in data.values():
+            self._set_graph(value)
+
+        super().__init__(data)
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        if value._graph is not None and value._graph is not self._graph:
+            raise ValueError(
+                f"Value '{value}' is already an initializer of a different graph. Please remove the value from the previous graph first"
+            )
+        value._is_initializer = True
+        value._graph = self._graph
+
+    def _maybe_unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        assert value._graph is self._graph, "Bug: value does not belong to the graph"
+        value._is_initializer = False
+        if value._owned_by_graph():
+            # Keep the graph reference if the value is still an input or an initializer
+            return
+        value._graph = None
+
+    def __setitem__(self, key: str, value: _core.Value) -> None:
+        """Set an initializer for the graph."""
+        if not isinstance(value, _core.Value):
+            raise TypeError(f"value must be a Value object, not {type(value)}")
+        if not isinstance(key, str):
+            raise TypeError(f"Value name must be a string, not {type(key)}")
+        if key == "":
+            raise ValueError("Value name cannot be an empty string")
+        if not value.name:
+            logger.info("Value %s does not have a name, setting it to '%s'", value, key)
+            value.name = key
+        elif key != value.name:
+            raise ValueError(
+                f"Key '{key}' does not match the name of the value '{value.name}'. Please use the value.name as the key."
+            )
+        if value.producer() is not None:
+            raise ValueError(
+                f"Value '{value}' is produced by a node and cannot be a graph initializer"
+            )
+        if key in self.data:
+            # If the key already exists, unset the old value
+            old_value = self.data[key]
+            self._maybe_unset_graph(old_value)
+        # Must call _set_graph before super().__setitem__ so that when there is an error,
+        # the dictionary is not modified
+        self._set_graph(value)
+        super().__setitem__(key, value)
+
+    def __delitem__(self, key: str) -> None:
+        """Delete an initializer from the graph."""
+        value = self.data[key]
+        # Must call _maybe_unset_graph before super().__delitem__ so that when there is an error,
+        # the dictionary is not modified
+        self._maybe_unset_graph(value)
+        super().__delitem__(key)
+
+    def add(self, value: _core.Value) -> None:
+        """Add an initializer to the graph."""
+        self[value.name] = value  # type: ignore[index]
+
+
+class Attributes(collections.UserDict[str, "_core.Attr"]):
+    """The attributes of a Node as ``dict[str, Attr]`` with additional access methods."""
+
+    def __init__(self, attrs: Iterable[_core.Attr]):
+        super().__init__({attr.name: attr for attr in attrs})
+
+    def __setitem__(self, key: str, value: _core.Attr) -> None:
+        """Set an attribute for the node."""
+        if type(key) is not str:
+            raise TypeError(f"Key must be a string, not {type(key)}")
+        if not isinstance(value, _core.Attr):
+            raise TypeError(f"Value must be an Attr, not {type(value)}")
+        super().__setitem__(key, value)
+
+    def add(self, value: _core.Attr) -> None:
+        """Add an attribute to the node."""
+        self[value.name] = value
+
+    def get_int(self, key: str, default: T = None) -> int | T:  # type: ignore[assignment]
+        """Get the integer value of the attribute."""
+        if key in self:
+            return self[key].as_int()
+        return default
+
+    def get_float(self, key: str, default: T = None) -> float | T:  # type: ignore[assignment]
+        """Get the float value of the attribute."""
+        if key in self:
+            return self[key].as_float()
+        return default
+
+    def get_string(self, key: str, default: T = None) -> str | T:  # type: ignore[assignment]
+        """Get the string value of the attribute."""
+        if key in self:
+            return self[key].as_string()
+        return default
+
+    def get_tensor(self, key: str, default: T = None) -> _protocols.TensorProtocol | T:  # type: ignore[assignment]
+        """Get the tensor value of the attribute."""
+        if key in self:
+            return self[key].as_tensor()
+        return default
+
+    def get_graph(self, key: str, default: T = None) -> _core.Graph | T:  # type: ignore[assignment]
+        """Get the graph value of the attribute."""
+        if key in self:
+            return self[key].as_graph()
+        return default
+
+    def get_ints(self, key: str, default: T = None) -> Sequence[int] | T:  # type: ignore[assignment]
+        """Get the Sequence of integers from the attribute."""
+        if key in self:
+            return self[key].as_ints()
+        return default
+
+    def get_floats(self, key: str, default: T = None) -> Sequence[float] | T:  # type: ignore[assignment]
+        """Get the Sequence of floats from the attribute."""
+        if key in self:
+            return self[key].as_floats()
+        return default
+
+    def get_strings(self, key: str, default: T = None) -> Sequence[str] | T:  # type: ignore[assignment]
+        """Get the Sequence of strings from the attribute."""
+        if key in self:
+            return self[key].as_strings()
+        return default
+
+    def get_tensors(
+        self,
+        key: str,
+        default: T = None,  # type: ignore[assignment]
+    ) -> Sequence[_protocols.TensorProtocol] | T:
+        """Get the Sequence of tensors from the attribute."""
+        if key in self:
+            return self[key].as_tensors()
+        return default
+
+    def get_graphs(self, key: str, default: T = None) -> Sequence[_core.Graph] | T:  # type: ignore[assignment]
+        """Get the Sequence of graphs from the attribute."""
+        if key in self:
+            return self[key].as_graphs()
+        return default

onnx_ir/_io.py

@@ -0,0 +1,133 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Load and save ONNX models."""
+
+from __future__ import annotations
+
+__all__ = ["load", "save"]
+
+import os
+from typing import Callable
+
+import onnx  # noqa: TID251
+
+from onnx_ir import _core, _protocols, serde
+from onnx_ir import external_data as _external_data
+from onnx_ir._polyfill import zip
+
+
+def load(path: str | os.PathLike, format: str | None = None) -> _core.Model:
+    """Load an ONNX model from a file.
+
+    Args:
+        path: The path to the ONNX file.
+        format: The format of the file (e.g. protobuf, textproto, json, etc.).
+            If None, the format is inferred from the file extension.
+
+    Returns:
+        The loaded model.
+    """
+    # Do not use ONNX to load external data because the IR handles external data
+    # by doing memory mapping directly.
+    proto = onnx.load(path, format=format, load_external_data=False)
+    model = serde.deserialize_model(proto)
+    base_dir = os.path.dirname(path)
+    # Set the base directory for external data to the directory of the ONNX file
+    # so that relative paths are resolved correctly.
+    _external_data.set_base_dir(model.graph, base_dir)
+    return model
+
+
+def save(
+    model: _core.Model,
+    path: str | os.PathLike,
+    format: str | None = None,
+    external_data: str | os.PathLike | None = None,
+    size_threshold_bytes: int = 256,
+    callback: Callable[[_protocols.TensorProtocol, _external_data.CallbackInfo], None]
+    | None = None,
+) -> None:
+    """Save an ONNX model to a file.
+
+    The model remains unchanged after the call. If any existing external tensor
+    references the provided ``external_data`` path, it will be invalidated
+    after the external data is overwritten. To obtain a valid model, use :func:`load`
+    to load the newly saved model, or provide a different external data path that
+    is not currently referenced by any tensors in the model.
+
+    .. tip::
+
+        A simple progress bar can be implemented by passing a callback function as the following::
+
+            import onnx_ir as ir
+            import tqdm
+
+            with tqdm.tqdm() as pbar:
+                total_set = False
+
+                def callback(tensor: ir.TensorProtocol, metadata: ir.external_data.CallbackInfo) -> None:
+                    nonlocal total_set
+                    if not total_set:
+                        pbar.total = metadata.total
+                        total_set = True
+
+                    pbar.update()
+                    pbar.set_description(f"Saving {tensor.name} ({tensor.dtype}, {tensor.shape}) at offset {metadata.offset}")
+
+                ir.save(
+                    ...,
+                    callback=callback,
+                )
+
+    Args:
+        model: The model to save.
+        path: The path to save the model to. E.g. "model.onnx".
+        format: The format of the file (e.g. ``protobuf``, ``textproto``, ``json``, etc.).
+            If None, the format is inferred from the file extension.
+        external_data: The relative path to save external data to. When specified,
+            all initializers in the model will be converted to external data and
+            saved to the specified directory. If None, all tensors will be saved unmodified.
+            That is, if a tensor in the model is already external, it will be saved
+            with the same external information; if the tensor is not external,
+            it will be serialized in the ONNX Proto message.
+        size_threshold_bytes: Save to external data if the tensor size in bytes is larger than this threshold.
+            Effective only when ``external_data`` is set.
+        callback: A callback function that is called for each tensor that is saved to external data
+            for debugging or logging purposes.
+
+    Raises:
+        ValueError: If the external data path is an absolute path.
+    """
+    if external_data is not None:
+        if os.path.isabs(external_data):
+            raise ValueError(
+                f"The external data path must be relative to the ONNX file path, not '{external_data}'."
+            )
+        base_dir = os.path.dirname(path)
+
+        # Store the original initializer values so they can be restored if modify_model=False
+        initialized_values: list[_core.Value] = []
+        for graph in model.graphs():
+            # Collect from all subgraphs as well
+            initialized_values.extend(graph.initializers.values())
+        tensors = [v.const_value for v in initialized_values]
+
+        try:
+            model = _external_data.unload_from_model(
+                model,
+                base_dir,
+                external_data,
+                size_threshold_bytes=size_threshold_bytes,
+                callback=callback,
+            )
+            proto = serde.serialize_model(model)
+            onnx.save(proto, path, format=format)
+
+        finally:
+            # Restore the original initializer values so the model is unchanged
+            for initializer, tensor in zip(initialized_values, tensors, strict=True):
+                initializer.const_value = tensor
+
+    else:
+        proto = serde.serialize_model(model)
+        onnx.save(proto, path, format=format)