[Bug Fix] Fix incorrect argument inference in AutoWrapper with self assignment (#1616)

SUMMARY: 
Fix a bug in the AutoWrapper `_wrap_stmt` logic where a variable that is
both read and written (e.g., `x = x + 1`) is incorrectly excluded from
the unbound argument list. This caused some internal variables (like
`attention_mask`) to be missing from wrapped function inputs, leading to
runtime errors when using torch.fx.

An example of reproducing this bug:
```python
import ast
from llmcompressor.pipelines.sequential.ast_utils.name_analyzer import NameAnalyzer
from llmcompressor.pipelines.sequential.ast_utils.auto_wrapper import AutoWrapper

code = """
def forward(x, y):
    if y > 0:
        x = x + 1
    else:
        x = x - 1
    return x
"""

tree = ast.parse(code)

analyzer = NameAnalyzer(omit=set())
unbound, assigned, cond = analyzer.analyze(tree)
print("=== Result ===")
print("Unbound:", unbound)
print("Assigned:", assigned)
print("Conditionally Assigned:", cond)

namespace = {}
wrapper = AutoWrapper(namespace, ignore=[])
wrapped_tree = wrapper.auto_wrap(ast.Module(body=[tree], type_ignores=[]))

print("\n=== Wrapped Code ===")
print(ast.unparse(wrapped_tree))
```

The output of the original code:
```
=== Result ===
Unbound: {'y'}
Assigned: {'x'}
Conditionally Assigned: set()

=== Wrapped Code ===
@torch.fx.wrap
def wrapped_8751792970731(y):
    if y > 0:
        x = x + 1
    else:
        x = x - 1
    return (x,)

def forward(x, y):
    x, = wrapped_8751792970731(y)
    return x
```
`x` is missing from the input args of `wrapped_8751792970731`, leading
to unbound variable in the wrapped function

The fix add a `visit_Assign` function to properly detect and classify
variables as unbound when they are read before assignment, even within
the same statement.

TEST PLAN: 
Added a new test case `test_branch_with_self_assignment` to
`tests/llmcompressor/pipelines/sequential/ast_utils.py/test_auto_wrapper.py`

---------

Signed-off-by: xiezhe.24 <xiezhe.24@bytedance.com>
Co-authored-by: xiezhe.24 <xiezhe.24@bytedance.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: 3 people

src/llmcompressor/pipelines/sequential/ast_utils/name_analyzer.py

@@ -66,6 +66,14 @@ def visit_Name(self, node: ast.Name) -> ast.Name:
 
         self.generic_visit(node)
 
+    def visit_Assign(self, node: ast.Assign):
+        # Visit the right side of the assignment first
+        self.visit(node.value)
+
+        # Now visit the left side of the assignment
+        for target in node.targets:
+            self.visit(target)
+
     def visit_If(self, node: ast.If):
         self.visit(node.test)
 

tests/llmcompressor/pipelines/sequential/ast_utils.py/test_auto_wrapper.py

@@ -98,3 +98,28 @@ def forward(self):
     """
     namespace = {"self": Model()}
     check_wrapping(source, None, 1, namespace=namespace, ignore=["meth_one"])
+
+
+def test_branch_with_self_assignment():
+    source = """
+    def forward(x, y):
+        if y > 0:
+            x = x + 1
+        else:
+            x = x - 1
+        return x
+    """
+
+    tree = ast.parse(textwrap.dedent(source))
+    wrapper = AutoWrapper(namespace={}, ignore=[])
+    wrapper.auto_wrap(tree)
+
+    assert len(wrapper._wrapper_fn_defs) == 1
+
+    # Check if both x, y are included in args
+    wrapped_fn = wrapper._wrapper_fn_defs[0]
+    arg_names = {arg.arg for arg in wrapped_fn.args.args}
+
+    assert arg_names == {"x", "y"}, (
+        f"Expected arguments {{'x', 'y'}}, but got {arg_names}"
+    )