Skip to content

[IR] Display constant tensors for Value #2248

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
May 1, 2025
2 changes: 1 addition & 1 deletion onnxscript/ir/_convenience/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -212,7 +212,7 @@ def convert_attributes(
... "type_protos": [ir.TensorType(ir.DataType.FLOAT), ir.TensorType(ir.DataType.FLOAT)],
... }
>>> convert_attributes(attrs)
[Attr('int', INT, 1), Attr('float', FLOAT, 1.0), Attr('str', STRING, 'hello'), Attr('ints', INTS, [1, 2, 3]), Attr('floats', FLOATS, [1.0, 2.0, 3.0]), Attr('strings', STRINGS, ['hello', 'world']), Attr('tensor', TENSOR, Tensor<DOUBLE,[3]>(array([1., 2., 3.]), name=None)), Attr('tensor_proto', TENSOR, TensorProtoTensor<FLOAT,[3]>(name='proto')), Attr('graph', INTS, Graph(
[Attr('int', INT, 1), Attr('float', FLOAT, 1.0), Attr('str', STRING, 'hello'), Attr('ints', INTS, [1, 2, 3]), Attr('floats', FLOATS, [1.0, 2.0, 3.0]), Attr('strings', STRINGS, ['hello', 'world']), Attr('tensor', TENSOR, Tensor<DOUBLE,[3]>(array([1., 2., 3.]), name=None)), Attr('tensor_proto', TENSOR, TensorProtoTensor<FLOAT,[3]>(array([1., 2., 3.], dtype=float32), name='proto')), Attr('graph', INTS, Graph(
name='graph0',
inputs=(
<BLANKLINE>
Expand Down
35 changes: 28 additions & 7 deletions onnxscript/ir/_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -406,7 +406,10 @@ def __dlpack_device__(self) -> tuple[int, int]:
return self.__array__().__dlpack_device__()

def __repr__(self) -> str:
return f"{self._repr_base()}({self._raw!r}, name={self.name!r})"
# Avoid multi-line repr
tensor_lines = repr(self._raw).split("\n")
tensor_text = " ".join(line.strip() for line in tensor_lines)
return f"{self._repr_base()}({tensor_text}, name={self.name!r})"

@property
def dtype(self) -> _enums.DataType:
Expand Down Expand Up @@ -1465,7 +1468,7 @@ def __str__(self) -> str:
+ ", ".join(
[
(
f"%{_quoted(x.name) if x.name else 'anonymous:' + str(id(x))}"
f"%{_quoted(x.name) if x.name else 'anonymous:' + str(id(x))}{x._constant_tensor_part()}"
if x is not None
else "None"
)
Expand Down Expand Up @@ -1836,14 +1839,20 @@ def __init__(

def __repr__(self) -> str:
value_name = self.name if self.name else "anonymous:" + str(id(self))
type_text = f", type={self.type!r}" if self.type is not None else ""
shape_text = f", shape={self.shape!r}" if self.shape is not None else ""
producer = self.producer()
if producer is None:
producer_text = "None"
producer_text = ""
elif producer.name is not None:
producer_text = producer.name
producer_text = f", producer='{producer.name}'"
else:
producer_text = f"anonymous_node:{id(producer)}"
return f"{self.__class__.__name__}({value_name!r}, type={self.type!r}, shape={self.shape}, producer={producer_text}, index={self.index()})"
producer_text = f", producer=anonymous_node:{id(producer)}"
index_text = f", index={self.index()}" if self.index() is not None else ""
const_value_text = self._constant_tensor_part()
if const_value_text:
const_value_text = f", const_value={const_value_text}"
return f"{self.__class__.__name__}(name={value_name!r}{type_text}{shape_text}{producer_text}{index_text}{const_value_text})"

def __str__(self) -> str:
value_name = self.name if self.name is not None else "anonymous:" + str(id(self))
Expand All @@ -1852,7 +1861,19 @@ def __str__(self) -> str:

# Quote the name because in reality the names can have invalid characters
# that make them hard to read
return f"%{_quoted(value_name)}<{type_text},{shape_text}>"
return (
f"%{_quoted(value_name)}<{type_text},{shape_text}>{self._constant_tensor_part()}"
)

def _constant_tensor_part(self) -> str:
"""Display string for the constant tensor attached to str of Value."""
if self.const_value is not None:
# Only display when the const value is small
if self.const_value.size <= 10:
return f"{{{self.const_value}}}"
else:
return f"{{{self.const_value.__class__.__name__}(...)}}"
return ""

def producer(self) -> Node | None:
"""The node that produces this value.
Expand Down
7 changes: 4 additions & 3 deletions onnxscript/ir/serde.py
Original file line number Diff line number Diff line change
Expand Up @@ -278,9 +278,10 @@ def raw(self) -> onnx.TensorProto:
return self._proto

def __repr__(self) -> str:
# It is a little hard to display the content when there can be types
# unsupported by numpy
# Preferably we should display some content when the tensor is small
if self.size <= 10:
tensor_lines = repr(self.numpy()).split("\n")
tensor_text = " ".join(line.strip() for line in tensor_lines)
return f"{self._repr_base()}({tensor_text}, name={self.name!r})"
return f"{self._repr_base()}(name={self.name!r})"

def __array__(self, dtype: Any = None) -> np.ndarray:
Expand Down
Loading