From 855fc27b4c7d2e7c03707131ef14378615f8cd89 Mon Sep 17 00:00:00 2001 From: Philipp Schaad Date: Wed, 22 Jan 2025 16:22:00 +0100 Subject: [PATCH] Fix incorrect offsets during assignments with broadcasts (#1875) --- dace/frontend/python/newast.py | 13 +++++++++++-- tests/numpy/split_test.py | 10 ++++++---- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/dace/frontend/python/newast.py b/dace/frontend/python/newast.py index 62023f1d1b..c9cfb698b4 100644 --- a/dace/frontend/python/newast.py +++ b/dace/frontend/python/newast.py @@ -2725,8 +2725,17 @@ def _add_assignment(self, fake_subset = dace.subsets.Range(missing_dimensions + op_dimensions) - # use this fake subset to calculate the offset - fake_subset.offset(squeezed, True) + # Use this fake subset to calculate the offset. Constant indices are ignored, as they do not depend + # on the broadcasting operation. + offset_indices_to_ignore = set() + for i, idx in enumerate(inp_idx): + if not symbolic.issymbolic(pystr_to_symbolic(idx)): + offset_indices_to_ignore.add(i) + fake_subset_offs_indices = [] + for i in range(len(fake_subset)): + if i not in offset_indices_to_ignore: + fake_subset_offs_indices.append(i) + fake_subset.offset(squeezed, True, indices=fake_subset_offs_indices) # we access the inp subset using the computed offset # since the inp_subset may be missing leading dimensions, we reverse-zip-reverse diff --git a/tests/numpy/split_test.py b/tests/numpy/split_test.py index ef783768b3..69e6ecc4ea 100644 --- a/tests/numpy/split_test.py +++ b/tests/numpy/split_test.py @@ -127,7 +127,8 @@ def test_dsplit_4d(): return a, b, c -def test_compiletime_split(): +@pytest.mark.parametrize('out_idx', [0, 1]) +def test_compiletime_split(out_idx): @dace.program def tester(x, y, in_indices: dace.compiletime, out_index: dace.compiletime): @@ -138,9 +139,9 @@ def tester(x, y, in_indices: dace.compiletime, out_index: dace.compiletime): x = np.random.rand(1000, 8) y = np.zeros_like(x) - tester(x, y, (1, 2, 3, 4, 5, 7), 0) + tester(x, y, (1, 2, 3, 4, 5, 7), out_idx) ref = np.zeros_like(y) - tester.f(x, ref, (1, 2, 3, 4, 5, 7), 0) + tester.f(x, ref, (1, 2, 3, 4, 5, 7), out_idx) assert np.allclose(y, ref) @@ -158,4 +159,5 @@ def tester(x, y, in_indices: dace.compiletime, out_index: dace.compiletime): test_vsplit() test_hsplit() test_dsplit_4d() - test_compiletime_split() + test_compiletime_split(0) + test_compiletime_split(1)