Merge pull request #1296 from spcl/attribute-replacements-new-state

New State for Attribute Replacements

dace/frontend/python/newast.py

@@ -4612,7 +4612,12 @@ def visit_Attribute(self, node: ast.Attribute):
         # Try to find sub-SDFG attribute
         func = oprepo.Replacements.get_attribute(type(arr), node.attr)
         if func is not None:
-            return func(self, self.sdfg, self.last_state, result)
+            # A new state is likely needed here, e.g., for transposition (ndarray.T)
+            self._add_state('%s_%d' % (type(node).__name__, node.lineno))
+            self.last_state.set_default_lineinfo(self.current_lineinfo)
+            result = func(self, self.sdfg, self.last_state, result)
+            self.last_state.set_default_lineinfo(None)
+            return result
 
         # Otherwise, try to find compile-time attribute (such as shape)
         try:

tests/numpy/attribute_test.py

@@ -33,6 +33,28 @@ def test_attribute_in_ranged_loop_symbolic():
     assert np.allclose(a, regression)
 
 
+def test_attribute_new_state():
+
+    N, F_in, F_out, heads = 2, 3, 4, 5
+
+    @dace.program
+    def fn(a: dace.float64[N, F_in], b: dace.float64[N, heads, F_out], c: dace.float64[heads * F_out, F_in]):
+        tmp = a.T @ np.reshape(b, (N, heads * F_out))
+        c[:] = tmp.T
+
+    rng = np.random.default_rng(42)
+
+    a = rng.random((N, F_in))
+    b = rng.random((N, heads, F_out))
+    c_expected = np.zeros((heads * F_out, F_in))
+    c = np.zeros((heads * F_out, F_in))
+
+    fn.f(a, b, c_expected)
+    fn(a, b, c)
+    assert np.allclose(c, c_expected)
+
+
 if __name__ == '__main__':
     test_attribute_in_ranged_loop()
     test_attribute_in_ranged_loop_symbolic()
+    test_attribute_new_state()