improve display

xadupre · xadupre · commit df4f45e20cd0 · 2024-11-12T14:46:25.000+01:00
diff --git a/onnx_array_api/array_api/_onnx_common.py b/onnx_array_api/array_api/_onnx_common.py
@@ -93,6 +93,8 @@ def asarray(
                     v = TEagerTensor(va)
     elif isinstance(a, float):
         v = TEagerTensor(np.array(a, dtype=np.float64))
+    elif isinstance(a, complex):
+        v = TEagerTensor(np.array(a, dtype=np.complex128))
     elif isinstance(a, bool):
         v = TEagerTensor(np.array(a, dtype=np.bool_))
     elif isinstance(a, str):
diff --git a/onnx_array_api/plotting/text_plot.py b/onnx_array_api/plotting/text_plot.py
@@ -865,10 +865,11 @@ def str_node(indent, node):
                 content = ""
             line_name_new[init.name] = len(rows)
             if init.doc_string:
-                rows.append(
+                t = (
                     f"init: name={init.name!r} type={_get_type(init)} "
-                    f"shape={_get_shape(init)}{content}         -- {init.doc_string}"
+                    f"shape={_get_shape(init)}{content}"
                 )
+                rows.append(f"{t}{' ' * max(0, 70 - len(t))}-- {init.doc_string}")
                 continue
             rows.append(
                 f"init: name={init.name!r} type={_get_type(init)} "
@@ -1117,10 +1118,11 @@ def onnx_text_plot_io(model, verbose=False, att_display=None):
     for init in model.initializer:
 
         if init.doc_string:
-            rows.append(
+            t = (
                 f"init: name={init.name!r} type={_get_type(init)} "
-                f"shape={_get_shape(init)}         -- {init.doc_string}"
+                f"shape={_get_shape(init)}"
             )
+            rows.append(f"{t}{' ' * max(0, 70 - len(t))}-- {init.doc_string}")
             continue
         rows.append(
             f"init: name={init.name!r} type={_get_type(init)} "
