Skip to content

Commit 2766661

Browse files
authored
Create short tensor str for nodes (#2262)
The display of constant value is now more concise. ``` >>> from onnxscript import ir >>> n = ir.node("Add", [ir.Value(name="a", const_value=ir.tensor(1)), ir.Value(name="b", const_value=ir.tensor([1]*10))], name="n0") >>> print(n) %"anonymous:123273301338960"<?,?> ⬅️ ::Add(%"a"{1}, %"b"{[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}) >>> print(repr(n)) Node(name='n0', domain='', op_type='Add', inputs=(Value(name='a', const_value={Tensor<INT64,[]>(array(1), name=None)}), Value(name='b', const_value={Tensor<INT64,[10]>(array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1]), name=None)})), attributes=OrderedDict(), overload='', outputs=(Value(name='anonymous:135329937264592', producer='n0', index=0),), version=None, doc_string=None) ```
1 parent f407d47 commit 2766661

File tree

1 file changed

+13
-1
lines changed

1 file changed

+13
-1
lines changed

onnxscript/ir/_core.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1263,6 +1263,18 @@ class Usage(NamedTuple):
12631263
idx: int
12641264

12651265

1266+
def _short_tensor_str_for_node(x: Value) -> str:
1267+
if x.const_value is None:
1268+
return ""
1269+
if x.const_value.size <= 10:
1270+
try:
1271+
data = x.const_value.numpy().tolist()
1272+
except Exception: # pylint: disable=broad-except
1273+
return "{...}"
1274+
return f"{{{data}}}"
1275+
return "{...}"
1276+
1277+
12661278
class Node(_protocols.NodeProtocol, _display.PrettyPrintable):
12671279
"""IR Node.
12681280
@@ -1427,7 +1439,7 @@ def __str__(self) -> str:
14271439
+ ", ".join(
14281440
[
14291441
(
1430-
f"%{_quoted(x.name) if x.name else 'anonymous:' + str(id(x))}{x._constant_tensor_part()}"
1442+
f"%{_quoted(x.name) if x.name else 'anonymous:' + str(id(x))}{_short_tensor_str_for_node(x)}"
14311443
if x is not None
14321444
else "None"
14331445
)

0 commit comments

Comments
 (0)