diff --git a/b_asic/architecture.py b/b_asic/architecture.py
index bdf1495293de6c6f8f12df755cbab1be5290c8ba..20fd1a3a5584b1836069f1e55dcc606cb42ed19b 100644
--- a/b_asic/architecture.py
+++ b/b_asic/architecture.py
@@ -4,6 +4,8 @@ B-ASIC architecture classes.
 from collections import defaultdict
 from typing import Dict, List, Optional, Set, Tuple, cast
 
+from graphviz import Digraph
+
 from b_asic.process import MemoryVariable, OperatorProcess, PlainMemoryVariable
 from b_asic.resources import ProcessCollection
 
@@ -33,6 +35,8 @@ class Resource:
             raise ValueError("Do not create Resource with empty ProcessCollection")
         self._collection = process_collection
         self._entity_name = entity_name
+        self._input_count = -1
+        self._output_count = -1
 
     def __repr__(self):
         return self._entity_name
@@ -64,6 +68,43 @@ class Resource:
             raise ValueError("Entity name must be set")
         raise NotImplementedError
 
+    def _repr_mimebundle_(self, include=None, exclude=None):
+        return self._digraph()._repr_mimebundle_(include=include, exclude=exclude)
+
+    def _repr_jpeg_(self):
+        return self._digraph()._repr_mimebundle_(include=["image/jpeg"])["image/jpeg"]
+
+    def _repr_png_(self):
+        return self._digraph()._repr_mimebundle_(include=["image/png"])["image/png"]
+
+    def _digraph(self) -> Digraph:
+        dg = Digraph(node_attr={'shape': 'record'})
+        dg.node(self._entity_name, self._struct_def())
+        return dg
+
+    @property
+    def input_count(self) -> int:
+        """Number of input ports."""
+        return self._input_count
+
+    @property
+    def output_count(self) -> int:
+        """Number of output ports."""
+        return self._output_count
+
+    def _struct_def(self) -> str:
+        inputs = [f"in{i}" for i in range(self._input_count)]
+        outputs = [f"out{i}" for i in range(self._output_count)]
+        ret = ""
+        if inputs:
+            instrs = [f"<{instr}> {instr}" for instr in inputs]
+            ret += f"{{{'|'.join(instrs)}}}|"
+        ret += f"{self._entity_name}"
+        if outputs:
+            outstrs = [f"<{outstr}> {outstr}" for outstr in outputs]
+            ret += f"|{{{'|'.join(outstrs)}}}"
+        return ret
+
 
 class ProcessingElement(Resource):
     """
@@ -100,11 +141,21 @@ class ProcessingElement(Resource):
         self._operation_type = op_type
         self._type_name = op_type.type_name()
         self._entity_name = entity_name
+        self._input_count = ops[0].input_count
+        self._output_count = ops[0].output_count
 
     @property
     def processes(self) -> Set[OperatorProcess]:
         return {cast(OperatorProcess, p) for p in self._collection}
 
+    def input_count(self) -> int:
+        """Return number of input ports."""
+        raise NotImplementedError()
+
+    def output_count(self) -> int:
+        """Return number of output ports."""
+        raise NotImplementedError()
+
 
 class Memory(Resource):
     """
@@ -278,7 +329,7 @@ class Architecture:
 
     def get_interconnects_for_pe(
         self, pe: ProcessingElement
-    ) -> Tuple[List[Dict[str, int]], List[Dict[str, int]]]:
+    ) -> Tuple[List[Dict[Resource, int]], List[Dict[Resource, int]]]:
         """
         Return lists of dictionaries with interconnect information for a
         ProcessingElement.
@@ -320,6 +371,35 @@ class Architecture:
         """
         self._entity_name = entity_name
 
+    def _repr_mimebundle_(self, include=None, exclude=None):
+        return self._digraph()._repr_mimebundle_(include=include, exclude=exclude)
+
+    def _repr_jpeg_(self):
+        return self._digraph()._repr_mimebundle_(include=["image/jpeg"])["image/jpeg"]
+
+    def _repr_png_(self):
+        return self._digraph()._repr_mimebundle_(include=["image/png"])["image/png"]
+
+    def _digraph(self) -> Digraph:
+        dg = Digraph(node_attr={'shape': 'record'})
+        for mem in self._memories:
+            dg.node(mem._entity_name, mem._struct_def())
+        for pe in self._processing_elements:
+            dg.node(pe._entity_name, pe._struct_def())
+        for pe in self._processing_elements:
+            inputs, outputs = self.get_interconnects_for_pe(pe)
+            for i, inp in enumerate(inputs):
+                for source, cnt in inp.items():
+                    dg.edge(
+                        source._entity_name, f"{pe._entity_name}:in{i}", label=f"{cnt}"
+                    )
+            for o, outp in enumerate(outputs):
+                for dest, cnt in outp.items():
+                    dg.edge(
+                        f"{pe._entity_name}:out{0}", dest._entity_name, label=f"{cnt}"
+                    )
+        return dg
+
     @property
     def memories(self) -> Set[Memory]:
         return self._memories
diff --git a/b_asic/signal_flow_graph.py b/b_asic/signal_flow_graph.py
index af6dd758c26008859cbb2606acbda9a4363fd9d9..49884f86de4b4b18711a635e3c18ad6cac6dcc2d 100644
--- a/b_asic/signal_flow_graph.py
+++ b/b_asic/signal_flow_graph.py
@@ -4,7 +4,7 @@ B-ASIC Signal Flow Graph Module.
 Contains the signal flow graph operation.
 """
 
-import itertools as it
+import itertools
 import re
 from collections import defaultdict, deque
 from io import StringIO
@@ -26,7 +26,6 @@ from typing import (
 )
 
 from graphviz import Digraph
-from matplotlib.axes import itertools
 
 from b_asic.graph_component import GraphComponent
 from b_asic.operation import (
@@ -963,7 +962,7 @@ class SFG(AbstractOperation):
         first_op = no_inputs_queue.popleft()
         visited = {first_op}
         p_queue = PriorityQueue()
-        p_queue_entry_num = it.count()
+        p_queue_entry_num = itertools.count()
         # Negative priority as max-heap popping is wanted
         p_queue.put((-first_op.output_count, -next(p_queue_entry_num), first_op))
 
@@ -1020,7 +1019,7 @@ class SFG(AbstractOperation):
 
                 # Else fetch operation with the lowest input count that is not zero
                 elif seen_but_not_visited_count > 0:
-                    for i in it.count(start=1):
+                    for i in itertools.count(start=1):
                         seen_inputs_queue = seen_with_inputs_dict[i]
                         if seen_inputs_queue:
                             new_op = seen_inputs_queue.popleft()
diff --git a/test/test_architecture.py b/test/test_architecture.py
index 4ff42db11f381c722a1cae96e1e9a2772b70a6fd..e83ecf4e307405a636d57dfcb1ba0cb87b17562c 100644
--- a/test/test_architecture.py
+++ b/test/test_architecture.py
@@ -84,7 +84,16 @@ def test_architecture(schedule_direct_form_iir_lp_filter: Schedule):
         for operation in chain(adders, const_mults, inputs, outputs)
     ]
     for i, pe in enumerate(processing_elements):
-        pe.set_entity_name(f"{pe._type_name.upper()}-{i}")
+        pe.set_entity_name(f"{pe._type_name.upper()}{i}")
+        if pe._type_name == 'add':
+            s = (
+                'digraph {\n\tnode [shape=record]\n\t'
+                + pe._entity_name
+                + ' [label="{<in0> in0|<in1> in1}|'
+                + pe._entity_name
+                + '|{<out0> out0}"]\n}'
+            )
+            assert pe._digraph().source in (s, s + '\n')
 
     # Extract zero-length memory variables
     direct_conn, mvs = mvs.split_on_length()
@@ -95,13 +104,16 @@ def test_architecture(schedule_direct_form_iir_lp_filter: Schedule):
     ]
     assert len(memories) == 1
     for i, memory in enumerate(memories):
-        memory.set_entity_name(f"mem-{i}")
+        memory.set_entity_name(f"MEM{i}")
+        s = 'digraph {\n\tnode [shape=record]\n\tMEM0 [label=MEM0]\n}'
+        assert memory._digraph().source in (s, s + '\n')
 
     # Create architecture from
     architecture = Architecture(
         set(processing_elements), set(memories), direct_interconnects=direct_conn
     )
 
+    # assert architecture._digraph().source == "foo"
     for pe in processing_elements:
         print(pe)
         for operation in pe._collection: