From ee6a13d5f481b3b67a37ee9d9e1e84fa93bcdfa6 Mon Sep 17 00:00:00 2001
From: Lukas Truemper <lukas.truemper@outlook.de>
Date: Mon, 30 Oct 2023 21:45:31 +0100
Subject: [PATCH] WCRToAugAssign: Added memlet propagation

---
 dace/transformation/dataflow/wcr_conversion.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/dace/transformation/dataflow/wcr_conversion.py b/dace/transformation/dataflow/wcr_conversion.py
index 85a4690912..7f4fbc654d 100644
--- a/dace/transformation/dataflow/wcr_conversion.py
+++ b/dace/transformation/dataflow/wcr_conversion.py
@@ -9,7 +9,7 @@
 from dace import SDFG, SDFGState
 from dace.sdfg.state import StateSubgraphView
 from dace.transformation import helpers
-from dace import propagate_memlet
+from dace.sdfg.propagation import propagate_memlets_state
 
 
 class AugAssignToWCR(transformation.SingleStateTransformation):
@@ -256,6 +256,7 @@ def apply(self, state: SDFGState, sdfg: SDFG):
 
         # Remove input node and connector
         state.remove_memlet_path(inedge)
+        propagate_memlets_state(sdfg, state)
 
         # If outedge leads to non-transient, and this is a nested SDFG,
         # propagate outwards