WIP: fix numpy while loop

GridTools · Oct 28, 2024 · 4dadcf4 · 4dadcf4
1 parent 1a8aab8
commit 4dadcf4
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/src/gt4py/cartesian/gtc/numpy/oir_to_npir.py b/src/gt4py/cartesian/gtc/numpy/oir_to_npir.py
@@ -157,13 +157,12 @@ def visit_AssignStmt(
     def visit_While(
         self, node: oir.While, *, mask: Optional[npir.Expr] = None, **kwargs: Any
     ) -> npir.While:
-        cond = self.visit(node.cond, mask=mask, **kwargs)
+        cond_expr = self.visit(node.cond, **kwargs)
         if mask:
-            mask = npir.VectorLogic(op=common.LogicalOperator.AND, left=mask, right=cond)
-        else:
-            mask = cond
+            cond_expr = npir.VectorLogic(op=common.LogicalOperator.AND, left=mask, right=cond_expr)
+
         return npir.While(
-            cond=cond, body=utils.flatten_list(self.visit(node.body, mask=mask, **kwargs))
+            cond=cond_expr, body=utils.flatten_list(self.visit(node.body, mask=cond_expr, **kwargs))
         )
 
     def visit_HorizontalRestriction(