Fix wp.static() capturing global variables instead of loop variables #1167
Conversation
📝 WalkthroughWalkthroughTracks loop-local variable names during codegen so wp.static() inside for-loops no longer captures shadowed globals; adds tests exercising simple, expression, nested, and shadowed-variable loop scenarios and a changelog entry. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
📜 Recent review detailsConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Pro 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (1)
🧰 Additional context used🧠 Learnings (1)📚 Learning: 2026-01-13T06:46:05.869ZApplied to files:
🪛 Ruff (0.14.11)warp/tests/test_static.py710-710: Unused function argument: (ARG001) 776-776: (F405) ⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
🔇 Additional comments (6)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
warp/_src/codegen.py (1)
3744-3753: Loop variable tracking logic is correct.The implementation properly adds and removes loop variables when entering/exiting
Fornode scope. The use ofdiscard()instead ofremove()is a good defensive choice.Optional: Consider handling tuple unpacking in loop targets.
While the comment notes that tuple unpacking is rare, the current implementation silently ignores patterns like
for x, y in .... Consider either:
- Extracting all names from tuple targets and tracking them
- Raising a warning if tuple unpacking is used with
wp.static()This could prevent subtle bugs if users do use tuple unpacking.
warp/tests/test_static.py (1)
691-716: Test logic is correct, but parameter may be unused.The test function correctly validates that
wp.static()captures loop variables instead of shadowed globals. The expected outputs are accurate:
- Simple loop:
[0, 1, 2]from iteration values- Expression:
[1, 3, 5]fromi*2+1- Nested:
[0, 1, 10, 11]fromi*10+jConsider addressing the unused
testparameter.Static analysis indicates the
testparameter is unused. If it's required by the test framework interface, consider prefixing it with an underscore (_test) to indicate it's intentionally unused. Otherwise, you may be able to use it for assertions instead of relying on NumPy'sassert_array_equal.♻️ Optional: Use test parameter for assertions
If the test framework provides assertion methods, you could use them instead:
- np.testing.assert_array_equal(results.numpy(), np.array([0, 1, 2]), err_msg="Simple loop variable test failed") + test.assertTrue(np.array_equal(results.numpy(), np.array([0, 1, 2])), "Simple loop variable test failed")Or if not needed, rename to
_testto indicate it's intentionally unused.
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
CHANGELOG.mdwarp/_src/codegen.pywarp/tests/test_static.py
🧰 Additional context used
🪛 Ruff (0.14.10)
warp/tests/test_static.py
691-691: Unused function argument: test
(ARG001)
748-748: add_function_test may be undefined, or defined from star imports
(F405)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: ci / test-warp (ubuntu-24.04, ubuntu-x86_64)
- GitHub Check: ci / test-warp (macos-latest, macos)
- GitHub Check: ci / test-warp (windows-2025, windows)
- GitHub Check: ci / test-warp (ubuntu-24.04-arm, ubuntu-aarch64)
- GitHub Check: ci / build-docs
- GitHub Check: cppcheck
🔇 Additional comments (6)
CHANGELOG.md (1)
18-19: Changelog entry looks good.The fix is documented clearly and accurately. The entry follows the established format, includes the GitHub issue reference, and provides sufficient context for users encountering this issue.
warp/_src/codegen.py (2)
3739-3742: LGTM!The initialization of
loop_varsas a set is appropriate for tracking unique loop variable names. The comment clearly explains the purpose of preventingwp.static()from capturing shadowed global variables.
3758-3765: Correct implementation of loop variable detection.The code properly:
- Walks the AST of the static expression to find all referenced names
- Checks intersection with tracked loop variables
- Defers evaluation by setting the
has_unresolved_static_expressionsflag- Continues traversal to allow later codegen-time evaluation
This ensures that loop constants are available when the static expression is eventually evaluated.
warp/tests/test_static.py (3)
662-665: LGTM!The global test variables are well-named and their purpose is clearly documented. The values (999, 888) are deliberately different from loop iteration values to make it obvious if
wp.static()incorrectly captures the globals instead of loop variables.
667-689: Comprehensive test coverage for loop variable handling.The three test kernels effectively cover:
- Simple loop variable capture
- Loop variables in complex expressions
- Nested loops with multiple loop variables
Each kernel has a clear docstring explaining its purpose, and the test cases progressively increase in complexity. This provides good validation that
wp.static()correctly handles loop variables in various scenarios.
748-753: LGTM!The test registration follows the same pattern as other tests in the file and correctly:
- Associates the test with the TestStatic class
- Uses a descriptive test name
- Passes the test function and devices parameter
The static analysis warning about
add_function_testbeing undefined is a false positive—it's imported via the star import fromwarp.tests.unittest_utilsat line 24.
Greptile SummaryThis PR fixes a bug where Key Changes:
Technical Correctness: Confidence Score: 5/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant User as User Code
participant AST as AST Transformer
participant Adj as Adjoint
participant Eval as Static Evaluator
participant Codegen as Code Generator
User->>AST: Define kernel with wp.static(loop_var)
AST->>AST: visit_For(): Track loop variable in loop_vars
AST->>AST: visit_Call(): Detect wp.static() referencing loop_var
AST->>Adj: Set has_unresolved_static_expressions = True
AST->>AST: Defer evaluation (don't replace with constant)
Note over AST,Adj: wp.static() left as call node in AST
Adj->>Codegen: emit_For(): Begin loop unrolling
loop For each iteration i
Codegen->>Adj: Add loop constant to symbols[loop_var]
Codegen->>Eval: Evaluate wp.static(loop_var) with loop constant
Eval->>Eval: Build evaluation context from adj.symbols
Eval->>Eval: Find loop_var constant in symbols
Eval-->>Codegen: Return correct loop iteration value
Codegen->>Codegen: Generate code for loop body with constant
end
Note over Codegen: Loop unrolled with correct constants
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def visit_For(self, node): | ||
| # Track loop variable while visiting loop body (simple names only; | ||
| # tuple unpacking like `for x, y in ...` is rare in Warp kernels) | ||
| var_name = node.target.id if isinstance(node.target, ast.Name) else None | ||
| if var_name: | ||
| self.loop_vars.add(var_name) | ||
| result = self.generic_visit(node) | ||
| if var_name: | ||
| self.loop_vars.discard(var_name) # Remove when exiting loop scope | ||
| return result |
There was a problem hiding this comment.
The loop variable is added to self.loop_vars before visiting child nodes, which incorrectly treats it as "in scope" when visiting the iterator expression (node.iter). In Python semantics, the iterator expression is evaluated before the loop variable is bound.
This causes incorrect behavior when wp.static() references a variable with the same name as the loop variable in the iterator expression:
n = 5
for n in range(wp.static(n)): # wp.static(n) should use global n=5, not defer
passThe current implementation would defer wp.static(n) because n is in loop_vars, but it should evaluate to 5 using the global scope.
To fix this, only track the loop variable when visiting the loop body and orelse clause:
| def visit_For(self, node): | |
| # Track loop variable while visiting loop body (simple names only; | |
| # tuple unpacking like `for x, y in ...` is rare in Warp kernels) | |
| var_name = node.target.id if isinstance(node.target, ast.Name) else None | |
| if var_name: | |
| self.loop_vars.add(var_name) | |
| result = self.generic_visit(node) | |
| if var_name: | |
| self.loop_vars.discard(var_name) # Remove when exiting loop scope | |
| return result | |
| def visit_For(self, node): | |
| # Track loop variable while visiting loop body (simple names only; | |
| # tuple unpacking like `for x, y in ...` is rare in Warp kernels) | |
| var_name = node.target.id if isinstance(node.target, ast.Name) else None | |
| # Visit the iterator expression first, without the loop variable in scope | |
| self.visit(node.iter) | |
| # Now add loop variable to scope and visit the body | |
| if var_name: | |
| self.loop_vars.add(var_name) | |
| for stmt in node.body: | |
| self.visit(stmt) | |
| for stmt in node.orelse: | |
| self.visit(stmt) | |
| if var_name: | |
| self.loop_vars.discard(var_name) | |
| return node |
There was a problem hiding this comment.
@greptileai Your observation about Python semantics is technically correct—the iterator expression is evaluated before the loop variable is bound. However, we verified this edge case works correctly in practice. When wp.static(n) in the iterator is deferred, the subsequent evaluation at codegen time still finds the global n in the evaluation context, producing the correct result. No code change is needed.
shi-eric
force-pushed
the
ershi/fix-static-loop-var
branch
from
b7e5a81 to
3db7004
Compare
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Fix all issues with AI agents
In @warp/tests/test_static.py:
- Around line 776-781: Ruff flags F405 for the bare usage of add_function_test
due to star-imports; fix by adding a local noqa to that call site — append a
trailing comment "# noqa: F405" to the add_function_test invocation for
test_static_loop_variable_not_shadowed_by_global so the single-line suppression
silences the warning without changing imports.
- Around line 674-680: The triple-quoted docstring inside the @wp.kernel
function static_loop_var_in_expr_kernel can end up as a runtime string literal
inside the kernel body and interfere with static-analysis; remove or move that
docstring out of the kernel (e.g., convert it to a regular comment above the
function or place the descriptive string at module level) so the kernel body
contains no top-level string literal, keeping the loop and the use of
wp.static(_global_test_idx * 2 + 1) unchanged.
- Around line 682-689: The kernel static_nested_loop_kernel contains a Python
docstring inside the kernel body which can be misinterpreted by the compiler;
remove the triple-quoted docstring from inside the kernel (or convert it to a
regular comment) and place any description outside the kernel function so the
loop variables (_global_test_idx, _global_test_j), idx calculation, and
wp.static usage remain unchanged and the test still expects [0, 1, 10, 11].
- Around line 710-743: The test function
test_static_loop_variable_not_shadowed_by_global has an unused parameter `test`
causing Ruff ARG001—rename it to `_test` or `_` (or add `# noqa: ARG001`) to
silence the linter; additionally, replace uses of
`np.testing.assert_array_equal(...)` with the file's existing
`assert_np_equal(...)` helper for consistency (update the four assertions: the
ones checking results, results2, results3, results4) while preserving the same
expected arrays and messages.
- Around line 667-672: The kernel function static_loop_var_kernel currently
contains a leading docstring inside the @wp.kernel function which can become an
ast.Expr and break Warp codegen; remove the triple-quoted string literal from
inside the function and instead place the explanatory text as a regular comment
immediately above the def line (or convert it to inline comments within the
body), ensuring the function body only contains executable statements that use
results and wp.static(_global_test_idx) with the loop variable _global_test_idx.
🧹 Nitpick comments (1)
warp/tests/test_static.py (1)
662-665: Make the “shadowed global” intent clearer and safer from collisions.These globals are fine for reproducing the bug, but consider making them even more unique (e.g., include
static_loop_shadow_) to reduce any chance of name collision with other tests/modules.
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
CHANGELOG.mdwarp/_src/codegen.pywarp/tests/test_static.py
🚧 Files skipped from review as they are similar to previous changes (2)
- CHANGELOG.md
- warp/_src/codegen.py
🧰 Additional context used
🪛 Ruff (0.14.10)
warp/tests/test_static.py
710-710: Unused function argument: test
(ARG001)
776-776: add_function_test may be undefined, or defined from star imports
(F405)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: ci / build-warp (windows-2025, windows, ["cudart", "nvcc", "nvrtc_dev", "nvjitlink"])
🔇 Additional comments (1)
warp/tests/test_static.py (1)
691-708: Test correctly asserts the intended behavior; docstring usage is valid.The test properly validates that
wp.static()captures the loop variable's final value (2) after the inner loop completes, not the global value (999). The assertionresults == [2, 2]is correct: the outer loop executes twice, and each iterationwp.static(_global_test_idx)captures 2, matching Python's variable shadowing semantics.Regarding docstrings in kernels: docstrings are explicitly allowed in Warp kernels per the framework documentation ("Only pass, ellipsis (...), comments, or docstrings are allowed"). This kernel follows the same pattern as other kernels in the same test file (
static_loop_var_kernel,static_nested_loop_kernel, etc.), all with docstrings, making the approach consistent and correct.Likely an incorrect or invalid review comment.
| @wp.kernel | ||
| def static_loop_var_kernel(results: wp.array(dtype=int)): | ||
| """Kernel where wp.static() should capture the loop variable, not the global.""" | ||
| for _global_test_idx in range(3): | ||
| results[_global_test_idx] = wp.static(_global_test_idx) | ||
|
|
There was a problem hiding this comment.
Potential Warp compile hazard: docstring inside @wp.kernel.
If Warp’s AST/codegen doesn’t explicitly ignore function docstrings, the leading string literal becomes an ast.Expr statement and can break compilation. Safer to replace kernel docstrings with comments above the function.
Proposed fix
@wp.kernel
def static_loop_var_kernel(results: wp.array(dtype=int)):
- """Kernel where wp.static() should capture the loop variable, not the global."""
+ # wp.static() should capture the loop variable, not the global.
for _global_test_idx in range(3):
results[_global_test_idx] = wp.static(_global_test_idx)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @wp.kernel | |
| def static_loop_var_kernel(results: wp.array(dtype=int)): | |
| """Kernel where wp.static() should capture the loop variable, not the global.""" | |
| for _global_test_idx in range(3): | |
| results[_global_test_idx] = wp.static(_global_test_idx) | |
| @wp.kernel | |
| def static_loop_var_kernel(results: wp.array(dtype=int)): | |
| # wp.static() should capture the loop variable, not the global. | |
| for _global_test_idx in range(3): | |
| results[_global_test_idx] = wp.static(_global_test_idx) |
🤖 Prompt for AI Agents
In @warp/tests/test_static.py around lines 667 - 672, The kernel function
static_loop_var_kernel currently contains a leading docstring inside the
@wp.kernel function which can become an ast.Expr and break Warp codegen; remove
the triple-quoted string literal from inside the function and instead place the
explanatory text as a regular comment immediately above the def line (or convert
it to inline comments within the body), ensuring the function body only contains
executable statements that use results and wp.static(_global_test_idx) with the
loop variable _global_test_idx.
| @wp.kernel | ||
| def static_loop_var_in_expr_kernel(results: wp.array(dtype=int)): | ||
| """Kernel where loop variable is used in an arithmetic expression.""" | ||
| for _global_test_idx in range(3): | ||
| # Even in complex expressions, the loop variable should be used | ||
| results[_global_test_idx] = wp.static(_global_test_idx * 2 + 1) | ||
|
|
There was a problem hiding this comment.
Same docstring-in-kernel risk as above.
Also, nice coverage for “loop var used in expression” (this matches the reported bug shape well).
Proposed fix
@wp.kernel
def static_loop_var_in_expr_kernel(results: wp.array(dtype=int)):
- """Kernel where loop variable is used in an arithmetic expression."""
+ # Loop variable is used in an arithmetic expression.
for _global_test_idx in range(3):
# Even in complex expressions, the loop variable should be used
results[_global_test_idx] = wp.static(_global_test_idx * 2 + 1)🤖 Prompt for AI Agents
In @warp/tests/test_static.py around lines 674 - 680, The triple-quoted
docstring inside the @wp.kernel function static_loop_var_in_expr_kernel can end
up as a runtime string literal inside the kernel body and interfere with
static-analysis; remove or move that docstring out of the kernel (e.g., convert
it to a regular comment above the function or place the descriptive string at
module level) so the kernel body contains no top-level string literal, keeping
the loop and the use of wp.static(_global_test_idx * 2 + 1) unchanged.
| @wp.kernel | ||
| def static_nested_loop_kernel(results: wp.array(dtype=int)): | ||
| """Kernel with nested loops - both loop variables should be protected.""" | ||
| for _global_test_idx in range(2): | ||
| for _global_test_j in range(2): | ||
| idx = _global_test_idx * 2 + _global_test_j | ||
| results[idx] = wp.static(_global_test_idx * 10 + _global_test_j) | ||
|
|
There was a problem hiding this comment.
Same docstring-in-kernel risk; nested-loop coverage looks solid.
The idx mapping and expected values in the test look consistent ([0, 1, 10, 11]).
Proposed fix
@wp.kernel
def static_nested_loop_kernel(results: wp.array(dtype=int)):
- """Kernel with nested loops - both loop variables should be protected."""
+ # Nested loops: both loop variables should be protected from global capture.
for _global_test_idx in range(2):
for _global_test_j in range(2):
idx = _global_test_idx * 2 + _global_test_j
results[idx] = wp.static(_global_test_idx * 10 + _global_test_j)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @wp.kernel | |
| def static_nested_loop_kernel(results: wp.array(dtype=int)): | |
| """Kernel with nested loops - both loop variables should be protected.""" | |
| for _global_test_idx in range(2): | |
| for _global_test_j in range(2): | |
| idx = _global_test_idx * 2 + _global_test_j | |
| results[idx] = wp.static(_global_test_idx * 10 + _global_test_j) | |
| @wp.kernel | |
| def static_nested_loop_kernel(results: wp.array(dtype=int)): | |
| # Nested loops: both loop variables should be protected from global capture. | |
| for _global_test_idx in range(2): | |
| for _global_test_j in range(2): | |
| idx = _global_test_idx * 2 + _global_test_j | |
| results[idx] = wp.static(_global_test_idx * 10 + _global_test_j) |
🤖 Prompt for AI Agents
In @warp/tests/test_static.py around lines 682 - 689, The kernel
static_nested_loop_kernel contains a Python docstring inside the kernel body
which can be misinterpreted by the compiler; remove the triple-quoted docstring
from inside the kernel (or convert it to a regular comment) and place any
description outside the kernel function so the loop variables (_global_test_idx,
_global_test_j), idx calculation, and wp.static usage remain unchanged and the
test still expects [0, 1, 10, 11].
| def test_static_loop_variable_not_shadowed_by_global(test, device): | ||
| """Test that wp.static() inside a for loop correctly captures the loop variable. | ||
| When a global Python variable exists with the same name as a kernel loop variable, | ||
| wp.static() should use the loop variable's compile-time constant value (0, 1, 2, ...), | ||
| not the unrelated global variable. This prevents confusing behavior where the | ||
| presence of a global variable silently changes the kernel's output. | ||
| """ | ||
| with wp.ScopedDevice(device): | ||
| # Test 1: Simple loop variable | ||
| results = wp.zeros(3, dtype=int) | ||
| wp.launch(static_loop_var_kernel, dim=1, inputs=[results]) | ||
| np.testing.assert_array_equal(results.numpy(), np.array([0, 1, 2]), err_msg="Simple loop variable test failed") | ||
|
|
||
| # Test 2: Loop variable in arithmetic expression | ||
| results2 = wp.zeros(3, dtype=int) | ||
| wp.launch(static_loop_var_in_expr_kernel, dim=1, inputs=[results2]) | ||
| np.testing.assert_array_equal( | ||
| results2.numpy(), np.array([1, 3, 5]), err_msg="Loop variable in expression test failed" | ||
| ) | ||
|
|
||
| # Test 3: Nested loops - both variables protected | ||
| results3 = wp.zeros(4, dtype=int) | ||
| wp.launch(static_nested_loop_kernel, dim=1, inputs=[results3]) | ||
| np.testing.assert_array_equal(results3.numpy(), np.array([0, 1, 10, 11]), err_msg="Nested loop test failed") | ||
|
|
||
| # Test 4: Nested loops reusing the same variable name | ||
| # Tests counter-based tracking: global should not be captured even with shadowing | ||
| results4 = wp.zeros(2, dtype=int) | ||
| wp.launch(static_nested_loop_same_var_kernel, dim=1, inputs=[results4]) | ||
| # Per Python semantics: inner loop shadows outer, final value (2) persists | ||
| np.testing.assert_array_equal( | ||
| results4.numpy(), np.array([2, 2]), err_msg="Nested loop with same variable test failed" | ||
| ) |
There was a problem hiding this comment.
Fix Ruff ARG001 (test unused) and consider using existing assert_np_equal helper.
Ruff flags test as unused here (ARG001). Either rename it to _test / _ or add a # noqa: ARG001. Also, for consistency with the rest of this file, you may want to use assert_np_equal(...) instead of np.testing.assert_array_equal(...).
Proposed fix (minimal)
-def test_static_loop_variable_not_shadowed_by_global(test, device):
+def test_static_loop_variable_not_shadowed_by_global(_test, device):
"""Test that wp.static() inside a for loop correctly captures the loop variable.
@@
with wp.ScopedDevice(device):
@@
- np.testing.assert_array_equal(results.numpy(), np.array([0, 1, 2]), err_msg="Simple loop variable test failed")
+ assert_np_equal(results.numpy(), np.array([0, 1, 2]), tol=0.0)
@@
- np.testing.assert_array_equal(
- results2.numpy(), np.array([1, 3, 5]), err_msg="Loop variable in expression test failed"
- )
+ assert_np_equal(results2.numpy(), np.array([1, 3, 5]), tol=0.0)
@@
- np.testing.assert_array_equal(results3.numpy(), np.array([0, 1, 10, 11]), err_msg="Nested loop test failed")
+ assert_np_equal(results3.numpy(), np.array([0, 1, 10, 11]), tol=0.0)
@@
- np.testing.assert_array_equal(
- results4.numpy(), np.array([2, 2]), err_msg="Nested loop with same variable test failed"
- )
+ assert_np_equal(results4.numpy(), np.array([2, 2]), tol=0.0)🧰 Tools
🪛 Ruff (0.14.10)
710-710: Unused function argument: test
(ARG001)
🤖 Prompt for AI Agents
In @warp/tests/test_static.py around lines 710 - 743, The test function
test_static_loop_variable_not_shadowed_by_global has an unused parameter `test`
causing Ruff ARG001—rename it to `_test` or `_` (or add `# noqa: ARG001`) to
silence the linter; additionally, replace uses of
`np.testing.assert_array_equal(...)` with the file's existing
`assert_np_equal(...)` helper for consistency (update the four assertions: the
ones checking results, results2, results3, results4) while preserving the same
expected arrays and messages.
There was a problem hiding this comment.
If the pre-commit checks passed on this MR (which it did) then there are no issues with the current set of Ruff rules used in this project.
There was a problem hiding this comment.
@shi-eric, you're absolutely right—if the pre-commit checks passed, the code is compliant with the project's configured Ruff rules. I apologize for the confusion. The project's Ruff configuration likely doesn't include the ARG001 rule or has it configured differently than I indicated. Thank you for clarifying!
| add_function_test( | ||
| TestStatic, | ||
| "test_static_loop_variable_not_shadowed_by_global", | ||
| test_static_loop_variable_not_shadowed_by_global, | ||
| devices=devices, | ||
| ) |
There was a problem hiding this comment.
Address Ruff F405 for add_function_test on this new call site.
Since this file relies on from warp.tests.unittest_utils import *, Ruff flags this new usage (F405). Easiest is a local # noqa: F405 on the new call (or switch to explicit import, but that’s broader than this change).
Proposed fix
-add_function_test(
+add_function_test( # noqa: F405
TestStatic,
"test_static_loop_variable_not_shadowed_by_global",
test_static_loop_variable_not_shadowed_by_global,
devices=devices,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| add_function_test( | |
| TestStatic, | |
| "test_static_loop_variable_not_shadowed_by_global", | |
| test_static_loop_variable_not_shadowed_by_global, | |
| devices=devices, | |
| ) | |
| add_function_test( # noqa: F405 | |
| TestStatic, | |
| "test_static_loop_variable_not_shadowed_by_global", | |
| test_static_loop_variable_not_shadowed_by_global, | |
| devices=devices, | |
| ) |
🧰 Tools
🪛 Ruff (0.14.10)
776-776: add_function_test may be undefined, or defined from star imports
(F405)
🤖 Prompt for AI Agents
In @warp/tests/test_static.py around lines 776 - 781, Ruff flags F405 for the
bare usage of add_function_test due to star-imports; fix by adding a local noqa
to that call site — append a trailing comment "# noqa: F405" to the
add_function_test invocation for
test_static_loop_variable_not_shadowed_by_global so the single-line suppression
silences the warning without changing imports.
There was a problem hiding this comment.
If the pre-commit checks passed on this MR (which it did) then there are no issues with the current set of Ruff rules used in this project.
There was a problem hiding this comment.
@shi-eric, thank you for the clarification! You're absolutely right—if the pre-commit checks passed with your project's Ruff configuration, then the F405 warning isn't applicable here. I appreciate you pointing that out, and I'll note that your project's linting rules are already properly configured for this codebase.
✏️ Learnings added
Learnt from: shi-eric
Repo: NVIDIA/warp PR: 1167
File: warp/tests/test_static.py:776-781
Timestamp: 2026-01-13T06:46:05.869Z
Learning: In the NVIDIA/warp repository, star imports like `from warp.tests.unittest_utils import *` in test files are acceptable per the project's Ruff configuration, and F405 warnings should not be flagged if pre-commit checks pass.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
| def visit_For(self, node): | ||
| # Track loop variable while visiting loop body (simple names only; | ||
| # tuple unpacking like `for x, y in ...` is rare in Warp kernels) | ||
| var_name = node.target.id if isinstance(node.target, ast.Name) else None | ||
| if var_name: | ||
| self.loop_vars[var_name] = self.loop_vars.get(var_name, 0) + 1 | ||
| result = self.generic_visit(node) | ||
| if var_name: | ||
| self.loop_vars[var_name] -= 1 | ||
| if self.loop_vars[var_name] == 0: | ||
| del self.loop_vars[var_name] | ||
| return result |
There was a problem hiding this comment.
The loop variable is added to self.loop_vars before generic_visit(node) is called. This means when the iterator expression (node.iter) is visited, the loop variable is already tracked, which violates Python scoping rules.
In Python, the loop variable is NOT in scope when evaluating the iterator:
i = 10
for i in range(wp.static(i)): # wp.static(i) should use outer i=10, not defer
passWith the current implementation, wp.static(i) in the iterator would be incorrectly deferred because i is already in self.loop_vars.
Fix: Only visit the loop body with the variable tracked, not the entire For node. Replace generic_visit(node) with explicit visiting:
| def visit_For(self, node): | |
| # Track loop variable while visiting loop body (simple names only; | |
| # tuple unpacking like `for x, y in ...` is rare in Warp kernels) | |
| var_name = node.target.id if isinstance(node.target, ast.Name) else None | |
| if var_name: | |
| self.loop_vars[var_name] = self.loop_vars.get(var_name, 0) + 1 | |
| result = self.generic_visit(node) | |
| if var_name: | |
| self.loop_vars[var_name] -= 1 | |
| if self.loop_vars[var_name] == 0: | |
| del self.loop_vars[var_name] | |
| return result | |
| def visit_For(self, node): | |
| # Track loop variable while visiting loop body (simple names only; | |
| # tuple unpacking like `for x, y in ...` is rare in Warp kernels) | |
| var_name = node.target.id if isinstance(node.target, ast.Name) else None | |
| # Visit target and iterator BEFORE adding loop var to tracking | |
| node.target = self.visit(node.target) | |
| node.iter = self.visit(node.iter) | |
| # NOW add loop variable for body traversal | |
| if var_name: | |
| self.loop_vars[var_name] = self.loop_vars.get(var_name, 0) + 1 | |
| # Visit body with loop variable tracked | |
| node.body = [self.visit(stmt) for stmt in node.body] | |
| node.orelse = [self.visit(stmt) for stmt in node.orelse] | |
| # Clean up after visiting body | |
| if var_name: | |
| self.loop_vars[var_name] -= 1 | |
| if self.loop_vars[var_name] == 0: | |
| del self.loop_vars[var_name] | |
| return node |
When a global Python variable had the same name as a kernel for-loop variable, wp.static() incorrectly used the global value. The fix tracks loop variables during AST traversal and defers wp.static() evaluation when the expression references a loop variable. Fixes NVIDIAGH-1139 Signed-off-by: Eric Shi <[email protected]>
shi-eric
force-pushed
the
ershi/fix-static-loop-var
branch
from
3db7004 to
11764fd
Compare
1 participant
When a global Python variable had the same name as a kernel for-loop variable, wp.static() incorrectly used the global value. The fix tracks loop variables during AST traversal and defers wp.static() evaluation when the expression references a loop variable.
Fixes GH-1139
Description
Before your PR is "Ready for review"
__init__.pyi,docs/api_reference/,docs/language_reference/)pre-commit run -aSummary by CodeRabbit
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.