diff --git a/src/pystencils/backend/ast/structural.py b/src/pystencils/backend/ast/structural.py index b0dace4e0cf42f21fcb335a58d3e7150185defd2..5c8fca9adf9546683f3f7489eb137a00f1523b9d 100644 --- a/src/pystencils/backend/ast/structural.py +++ b/src/pystencils/backend/ast/structural.py @@ -247,7 +247,7 @@ class PsLoop(PsStructuralNode): self._start.clone(), self._stop.clone(), self._step.clone(), - self._body.clone(), + self._body._clone_structural(), ) def get_children(self) -> tuple[PsAstNode, ...]: @@ -312,8 +312,8 @@ class PsConditional(PsStructuralNode): def _clone_structural(self) -> PsConditional: return PsConditional( self._condition.clone(), - self._branch_true.clone(), - self._branch_false.clone() if self._branch_false is not None else None, + self._branch_true._clone_structural(), + self._branch_false._clone_structural() if self._branch_false is not None else None, ) def get_children(self) -> tuple[PsAstNode, ...]: diff --git a/src/pystencils/backend/transformations/loop_vectorizer.py b/src/pystencils/backend/transformations/loop_vectorizer.py index e1e4fea502c08de86e13de5e3c251f1b7a7d0ee6..6b518a30de31acac67c061ead1683c1d6ab06816 100644 --- a/src/pystencils/backend/transformations/loop_vectorizer.py +++ b/src/pystencils/backend/transformations/loop_vectorizer.py @@ -213,7 +213,7 @@ class LoopVectorizer: trailing_ctr = self._ctx.duplicate_symbol(scalar_ctr) trailing_loop_body = substitute_symbols( - loop.body.clone(), {scalar_ctr: PsExpression.make(trailing_ctr)} + loop.body._clone_structural(), {scalar_ctr: PsExpression.make(trailing_ctr)} ) trailing_loop = PsLoop( PsExpression.make(trailing_ctr),