onnx-ir 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

onnx_ir/__init__.py

@@ -14,6 +14,7 @@ __all__ = [
     "ExternalTensor",
     "StringTensor",
     "LazyTensor",
+    "PackedTensor",
     "SymbolicDim",
     "Shape",
     "TensorType",
@@ -73,6 +74,7 @@ __all__ = [
     "from_proto",
     "from_onnx_text",
     "to_proto",
+    "to_onnx_text",
     # Convenience constructors
     "tensor",
     "node",
@@ -114,6 +116,7 @@ from onnx_ir._core import (
     Model,
     Node,
     OptionalType,
+    PackedTensor,
     RefAttr,
     SequenceType,
     Shape,
@@ -149,7 +152,7 @@ from onnx_ir._protocols import (
     TypeProtocol,
     ValueProtocol,
 )
-from onnx_ir.serde import TensorProtoTensor, from_onnx_text, from_proto, to_proto
+from onnx_ir.serde import TensorProtoTensor, from_onnx_text, from_proto, to_onnx_text, to_proto
 
 DEBUG = False
 
@@ -164,4 +167,4 @@ def __set_module() -> None:
 
 
 __set_module()
-__version__ = "0.1.1"
+__version__ = "0.1.3"

onnx_ir/_convenience/__init__.py

@@ -14,14 +14,17 @@ __all__ = [
     "replace_all_uses_with",
     "create_value_mapping",
     "replace_nodes_and_values",
+    "get_const_tensor",
 ]
 
+import logging
 from collections.abc import Mapping, Sequence
 from typing import Union
 
-import onnx
+import numpy as np
+import onnx  # noqa: TID251
 
-from onnx_ir import _core, _enums, _protocols, serde
+from onnx_ir import _core, _enums, _protocols, serde, traversal
 
 SupportedAttrTypes = Union[
     str,
@@ -42,6 +45,9 @@ SupportedAttrTypes = Union[
 ]
 
 
+logger = logging.getLogger(__name__)
+
+
 def _infer_attribute_type(attr: SupportedAttrTypes) -> _enums.AttributeType:
     """Infer the attribute type based on the type of the Python object."""
     if isinstance(attr, int):
@@ -313,7 +319,12 @@ def replace_all_uses_with(
 def create_value_mapping(graph: _core.Graph) -> dict[str, _core.Value]:
     """Return a dictionary mapping names to values in the graph.
 
-    The mapping does not include values from subgraphs.
+    The mapping includes values from subgraphs. Duplicated names are omitted,
+    and the first value with that name is returned. Values with empty names
+    are excluded from the mapping.
+
+    .. versionchanged:: 0.1.2
+        Values from subgraphs are now included in the mapping.
 
     Args:
         graph: The graph to extract the mapping from.
@@ -327,11 +338,23 @@ def create_value_mapping(graph: _core.Graph) -> dict[str, _core.Value]:
     for input in graph.inputs:
         if not input.name:
             continue
+        if input.name in values:
+            continue
         values[input.name] = input
-    for node in graph:
+    for node in traversal.RecursiveGraphIterator(graph):
+        for value in node.inputs:
+            if not value:
+                continue
+            if not value.name:
+                continue
+            if value.name in values:
+                continue
+            values[value.name] = value
         for value in node.outputs:
             if not value.name:
                 continue
+            if value.name in values:
+                continue
             values[value.name] = value
     return values
 
@@ -375,3 +398,106 @@ def replace_nodes_and_values(
     # insert new nodes after the index node
     graph_or_function.insert_after(insertion_point, new_nodes)
     graph_or_function.remove(old_nodes, safe=True)
+
+
+def get_const_tensor(
+    value: _core.Value, propagate_shape_type: bool = False
+) -> _protocols.TensorProtocol | None:
+    """Get the constant tensor from a value, if it exists.
+
+    A constant tensor can be obtained if the value has a ``const_value`` set
+    (as in the case of an initializer) or if the value is produced by a
+    Constant node.
+
+    This function will not alter the ``const_value`` of the value, but
+    it will propagate the shape and type of the constant tensor to the value
+    if `propagate_shape_type` is set to True.
+
+    .. versionadded:: 0.1.2
+
+    Args:
+        value: The value to get the constant tensor from.
+        propagate_shape_type: If True, the shape and type of the value will be
+            propagated to the Value.
+
+    Returns:
+        The constant tensor if it exists, otherwise None.
+
+    Raises:
+        ValueError: If the Constant node does not have exactly one output or
+            one attribute.
+    """
+    tensor = None
+    if value.const_value is not None:
+        tensor = value.const_value
+    else:
+        node = value.producer()
+        if node is None:
+            # Potentially a graph input
+            return None
+        if node.op_type != "Constant" or node.domain != "":
+            # Not a Constant node or not in the ONNX domain
+            return None
+        if len(node.outputs) != 1:
+            raise ValueError(
+                f"Constant node '{node.name}' must have exactly one output, "
+                f"but has {len(node.outputs)} outputs."
+            )
+        if len(node.attributes) != 1:
+            raise ValueError(
+                f"Constant node '{node.name}' must have exactly one attribute, "
+                f"but has {len(node.attributes)} attributes."
+            )
+
+        attr_name, attr_value = next(iter(node.attributes.items()))
+
+        if attr_value.is_ref():
+            # TODO: Make it easier to resolve a reference attribute.
+            # For now we just return None
+            return None
+
+        ir_value = node.outputs[0]
+        if attr_name in {"value_float", "value_floats"}:
+            tensor = _core.Tensor(
+                np.array(attr_value.value, dtype=np.float32), name=ir_value.name
+            )
+        elif attr_name in {"value_int", "value_ints"}:
+            tensor = _core.Tensor(
+                np.array(attr_value.value, dtype=np.int64), name=ir_value.name
+            )
+        elif attr_name in {"value_string", "value_strings"}:
+            tensor = _core.StringTensor(
+                np.array(attr_value.value, dtype=np.bytes_), name=ir_value.name
+            )
+        elif attr_name == "value":
+            tensor = attr_value.as_tensor()
+        else:
+            raise ValueError(
+                f"Unsupported attribute '{attr_name}' in Constant node '{node.name}'. "
+                "Expected one of 'value_float', 'value_floats', 'value_int', "
+                "'value_ints', 'value_string', 'value_strings', or 'value'."
+            )
+        # Assign the name of the constant value to the tensor
+        tensor.name = value.name
+    if tensor is not None and propagate_shape_type:
+        # Propagate the shape and type of the tensor to the value
+        if value.shape is not None and value.shape != tensor.shape:
+            logger.warning(
+                "Value '%s' has a shape %s that differs from "
+                "the constant tensor's shape %s. The value's shape will be updated.",
+                value,
+                value.shape,
+                tensor.shape,
+            )
+        value.shape = tensor.shape  # type: ignore[assignment]
+        new_value_type = _core.TensorType(tensor.dtype)
+        if value.type is not None and value.type != new_value_type:
+            logger.warning(
+                "Value '%s' has a type '%s' that differs from "
+                "the constant tensor's type '%s'. The value's type will be updated.",
+                value,
+                value.type,
+                new_value_type,
+            )
+        value.type = new_value_type
+    return tensor

onnx_ir/_convenience/_constructors.py

@@ -13,7 +13,7 @@ import typing
 from collections.abc import Mapping, Sequence
 
 import numpy as np
-import onnx
+import onnx  # noqa: TID251
 
 from onnx_ir import _convenience, _core, _enums, _protocols, serde, tensor_adapters
 
@@ -37,6 +37,10 @@ def tensor(
 
     ``value`` can be a numpy array, a plain Python object, or a TensorProto.
 
+    .. warning::
+        For 4bit dtypes, the value must be unpacked. Use :class:`~onnx_ir.PackedTensor`
+        to create a tensor with packed data.
+
     Example::
 
         >>> import onnx_ir as ir
@@ -155,7 +159,7 @@ def node(
     doc_string: str | None = None,
     metadata_props: dict[str, str] | None = None,
 ) -> ir.Node:
-    """Create an :class:`ir.Node`.
+    """Create an :class:`~onnx_ir.Node`.
 
     This is a convenience constructor for creating a Node that supports Python
     objects as attributes.