Convenience function get_const_tensor (#45)

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.

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: codecov-ai[bot] <156709835+codecov-ai[bot]@users.noreply.github.com>

Author: justinchuby

docs/api/ir_convenience.md

@@ -9,7 +9,8 @@
 ```{eval-rst}
 .. autofunction:: convert_attribute
 .. autofunction:: convert_attributes
+.. autofunction:: create_value_mapping
+.. autofunction:: get_const_tensor
 .. autofunction:: replace_all_uses_with
 .. autofunction:: replace_nodes_and_values
-.. autofunction:: create_value_mapping
 ```

src/onnx_ir/_convenience/__init__.py

@@ -14,11 +14,14 @@
     "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 numpy as np
 import onnx  # noqa: TID251
 
 from onnx_ir import _core, _enums, _protocols, serde, traversal
@@ -42,6 +45,9 @@
 ]
 
 
+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):
@@ -389,3 +395,104 @@ 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.
+
+    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

src/onnx_ir/_convenience/_convenience_test.py

@@ -0,0 +1,103 @@
+# Copyright (c) ONNX Project Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+
+import numpy as np
+import parameterized
+
+import onnx_ir as ir
+from onnx_ir import _convenience
+
+
+class GetConstantTensorTest(unittest.TestCase):
+    def test_direct_const_value(self):
+        # Test when value has a direct const_value
+        tensor = ir.Tensor(np.array([1, 2, 3], dtype=np.int64), name="test_tensor")
+        value = ir.Value(name="test_value", type=ir.TensorType(ir.DataType.INT64))
+        value.const_value = tensor
+        self.assertIs(_convenience.get_const_tensor(value), tensor)
+
+    def test_no_const_value(self):
+        value = ir.Value(name="test_value", type=ir.TensorType(ir.DataType.FLOAT))
+
+        self.assertIsNone(_convenience.get_const_tensor(value))
+
+    def test_non_constant_producer_node(self):
+        # Test when producer node is not a Constant
+        node = ir.Node(
+            name="test_node",
+            domain="",
+            op_type="Add",
+            inputs=[],
+        )
+
+        output_value = node.outputs[0]
+        self.assertIsNone(_convenience.get_const_tensor(output_value))
+
+    @parameterized.parameterized.expand(
+        [
+            (
+                "value_float",
+                ir.AttrFloat32("value_float", 3.14),
+                np.array(3.14, dtype=np.float32),
+            ),
+            ("value_int", ir.AttrInt64("value_int", 42), np.array(42, dtype=np.int64)),
+            (
+                "value_string",
+                ir.AttrString("value_string", "test"),
+                np.array(b"test", dtype=object),
+            ),
+            (
+                "value_floats",
+                ir.AttrFloat32s("value_floats", [1.0, 2.0, 3.0]),
+                np.array([1.0, 2.0, 3.0], dtype=np.float32),
+            ),
+            (
+                "value_ints",
+                ir.AttrInt64s("value_ints", [1, 2, 3]),
+                np.array([1, 2, 3], dtype=np.int64),
+            ),
+            (
+                "value_strings",
+                ir.AttrStrings("value_strings", ["a", "b", "c"]),
+                np.array([b"a", b"b", b"c"], dtype=object),
+            ),
+            (
+                "value",
+                ir.AttrTensor("value", ir.tensor(np.array([1.0, 2.0, 3.0], dtype=np.float32))),
+                np.array([1.0, 2.0, 3.0], dtype=np.float32),
+            ),
+        ]
+    )
+    def test_constant_value(self, _: str, attr: ir.Attr, expected: np.ndarray):
+        # Test with Constant node with float value
+        node = ir.Node(
+            name="constant_node",
+            domain="",
+            op_type="Constant",
+            inputs=[],
+            attributes=(attr,),
+        )
+        node.outputs[0].name = "output"
+
+        result = _convenience.get_const_tensor(node.outputs[0])
+
+        self.assertIsNotNone(result)
+        self.assertEqual(result.name, "output")
+        np.testing.assert_array_equal(result.numpy(), expected)
+
+        self.assertIsNone(node.outputs[0].shape)
+        self.assertIsNone(node.outputs[0].type)
+
+        result_2 = _convenience.get_const_tensor(node.outputs[0], propagate_shape_type=True)
+        self.assertIsNotNone(result_2)
+        self.assertEqual(result_2.name, "output")
+        np.testing.assert_array_equal(result_2.numpy(), expected)
+        self.assertEqual(node.outputs[0].shape, expected.shape)
+        self.assertEqual(node.outputs[0].type, ir.TensorType(result_2.dtype))
+
+
+if __name__ == "__main__":
+    unittest.main()

src/onnx_ir/convenience.py

@@ -7,15 +7,17 @@
 __all__ = [
     "convert_attribute",
     "convert_attributes",
+    "create_value_mapping",
+    "get_const_tensor",
     "replace_all_uses_with",
     "replace_nodes_and_values",
-    "create_value_mapping",
 ]
 
 from onnx_ir._convenience import (
     convert_attribute,
     convert_attributes,
     create_value_mapping,
+    get_const_tensor,
     replace_all_uses_with,
     replace_nodes_and_values,
 )