diff --git a/b_asic/signal_flow_graph.py b/b_asic/signal_flow_graph.py
index 2be252af6210940b5921f77309394d5ea2452bb1..aa3c23f2816d7dc615f34720e9c0ca36e08dad33 100644
--- a/b_asic/signal_flow_graph.py
+++ b/b_asic/signal_flow_graph.py
@@ -1416,8 +1416,10 @@ class SFG(AbstractOperation):
                         destination.operation.graph_id,
                     )
             else:
-                if op.type_name() == Delay.type_name():
+                if isinstance(op, Delay):
                     dg.node(op.graph_id, shape="square")
+                elif isinstance(op, (Input, Output)):
+                    dg.node(op.graph_id, shape="cds")
                 else:
                     dg.node(op.graph_id)
         return dg
diff --git a/test/test_sfg.py b/test/test_sfg.py
index f4f9dc4a56b510761eed4657891903b013c1fb4e..6f69f1923495c97f653c189e25491fe281d8d5b5 100644
--- a/test/test_sfg.py
+++ b/test/test_sfg.py
@@ -1398,18 +1398,18 @@ class TestPrecedenceGraph:
 class TestSFGGraph:
     def test_sfg(self, sfg_simple_filter):
         res = (
-            "digraph {\n\trankdir=LR\n\tin1\n\tin1 -> "
-            "add1\n\tout1\n\tt1 -> out1\n\tadd1\n\tcmul1 -> "
+            "digraph {\n\trankdir=LR\n\tin1 [shape=cds]\n\tin1 -> "
+            "add1\n\tout1 [shape=cds]\n\tt1 -> out1\n\tadd1\n\tcmul1 -> "
             "add1\n\tcmul1\n\tadd1 -> t1\n\tt1 [shape=square]\n\tt1 "
             "-> cmul1\n}"
         )
-
+        print(sfg_simple_filter.sfg().source)
         assert sfg_simple_filter.sfg().source in (res, res + "\n")
 
     def test_sfg_show_id(self, sfg_simple_filter):
         res = (
-            "digraph {\n\trankdir=LR\n\tin1\n\tin1 -> add1 "
-            "[label=s1]\n\tout1\n\tt1 -> out1 [label=s2]\n\tadd1"
+            "digraph {\n\trankdir=LR\n\tin1 [shape=cds]\n\tin1 -> add1 "
+            "[label=s1]\n\tout1 [shape=cds]\n\tt1 -> out1 [label=s2]\n\tadd1"
             "\n\tcmul1 -> add1 [label=s3]\n\tcmul1\n\tadd1 -> t1 "
             "[label=s4]\n\tt1 [shape=square]\n\tt1 -> cmul1 [label=s5]\n}"
         )