Correct memlet dimensions in case of indirect indices

dace/frontend/python/newast.py

@@ -3493,6 +3493,10 @@ def _visit_assign(self, node, node_target, op, dtype=None, is_return=False):
             else:
                 new_name, new_rng = true_name, rng
 
+            # Change the range in case indirect indices are used
+            for dim, indarr in indirect_indices.items():
+                new_rng[dim] = (0, self.sdfg.arrays[indarr].shape[0] - 1, 1)
+
             # Self-copy check
             if result in self.views and new_name == self.views[result][1].data:
                 read_rng = self.views[result][1].subset

tests/python_frontend/augmented_assignment_to_slice_test.py

@@ -1,6 +1,8 @@
 # Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved.
 
 import dace
+import numpy as np
+
 
 def test_augmented_assignment_to_indirect_access():
 
@@ -14,6 +16,7 @@ def _test_prog(A: dace.int32[M], ind: dace.int32[N], B: dace.int32[N]):
     sdfg = _test_prog.to_sdfg()
     assert sdfg.is_valid()
 
+
 def test_augmented_assignment_to_indirect_access_regression():
 
     N = dace.symbol('N')
@@ -26,6 +29,16 @@ def _test_prog(A: dace.int32[M], ind: dace.int32[N], B: dace.int32[N]):
     sdfg = _test_prog.to_sdfg()
     assert sdfg.is_valid()
 
+    # Test correctness
+    A = np.random.randint(0, 100, 10).astype(np.int32)
+    # Create a random permutation so that we don't create duplicate indices (as per NumPy)
+    ind = np.random.permutation(10).astype(np.int32)[:5]
+    B = np.random.randint(0, 100, 5).astype(np.int32)
+    A_copy = A.copy()
+    A_copy[ind] += B
+    _test_prog(A, ind, B)
+    assert np.allclose(A, A_copy)
+
 
 if __name__ == '__main__':
     test_augmented_assignment_to_indirect_access()