diff --git a/litedram/frontend/axi.py b/litedram/frontend/axi.py
index 33eb758..19dbc75 100644
--- a/litedram/frontend/axi.py
+++ b/litedram/frontend/axi.py
@@ -82,7 +82,7 @@ class LiteDRAMAXIBurst2Beat(Module):
 
         # # #
 
-        count = Signal(8)
+        self.count = count = Signal(8)
         size = Signal(8 + 4)
         offset = Signal(8 + 4)
 
@@ -151,10 +151,11 @@ class LiteDRAMAXI2NativeW(Module):
 
         # Burst to Beat
         aw_buffer = stream.Buffer(ax_description(axi.address_width, axi.id_width))
+        self.submodules += aw_buffer
         self.comb += axi.aw.connect(aw_buffer.sink)
         aw = stream.Endpoint(ax_description(axi.address_width, axi.id_width))
         aw_burst2beat = LiteDRAMAXIBurst2Beat(aw_buffer.source, aw)
-        self.submodules += aw_buffer, aw_burst2beat
+        self.submodules.aw_burst2beat = aw_burst2beat
 
         # Write Buffer
         w_buffer = stream.SyncFIFO(w_description(axi.data_width), buffer_depth)
@@ -215,10 +216,11 @@ class LiteDRAMAXI2NativeR(Module):
 
         # Burst to Beat
         ar_buffer = stream.Buffer(ax_description(axi.address_width, axi.id_width))
+        self.submodules += ar_buffer
         self.comb += axi.ar.connect(ar_buffer.sink)
         ar = stream.Endpoint(ax_description(axi.address_width, axi.id_width))
         ar_burst2beat = LiteDRAMAXIBurst2Beat(ar_buffer.source, ar)
-        self.submodules += ar_buffer, ar_burst2beat
+        self.submodules.ar_burst2beat = ar_burst2beat
 
         # Read buffer
         r_buffer = stream.SyncFIFO(r_description(axi.data_width, axi.id_width), buffer_depth)