diff --git a/pystencils/astnodes.py b/pystencils/astnodes.py index 79222a05a35e51f5581c01f95609e9ffc153524a..f4879ffe808be903dec78d353008aa77b44b0511 100644 --- a/pystencils/astnodes.py +++ b/pystencils/astnodes.py @@ -622,7 +622,8 @@ class SympyAssignment(Node): result.update(self._lhs_symbol.atoms(sp.Symbol)) - sizes = set().union(*(a.field.shape for a in self._lhs_symbol.atoms(ResolvedFieldAccess))) + sizes = set().union(*(a.field.shape[:a.field.spatial_dimensions] + for a in self._lhs_symbol.atoms(ResolvedFieldAccess))) sizes = filter(lambda s: isinstance(s, FieldShapeSymbol), sizes) result.update(sizes)