Fix typecheck

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, ...]:

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),