[Autowrapper] Fix local names, increase reproducability

[kylesayrs]

Purpose

• Fix bug where args and kwargs were not being recognized as local names in the namespace when autowrapping
  • This was a problem for the Kimi K2 model which includes this line

def forward(..., kwargs):
    if "padding_mask" in kwargs:  # kwargs was not being recognized as a local name
        ...

• Make autowrapper easier to debug by making wrapped function names and arguments consistent

Changes

• Add positional arguments and kwargs to set of local names
• Add an incrementing counter which is used to name wrapped functions, rather than naming them with a hash
• Refactored autowrapper tests to be easier to follow

Testing

• Can trace K2 model with higher granularity
• Added test test_function_variadic

[gemini-code-assist]

Summary of Changes

Hello @kylesayrs, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly improves the Autowrapper utility by addressing a critical bug that prevented proper recognition of args and kwargs as local names, which was necessary for certain model architectures. Beyond the bug fix, the changes also focus on making the autowrapper's behavior more deterministic and easier to debug by standardizing the naming of wrapped functions and the ordering of their arguments. The accompanying test suite has been updated and expanded to reflect these improvements and ensure robust functionality.

Highlights

• Bug Fix: Local Name Recognition: Resolved a critical bug where args and kwargs were not correctly recognized as local names within the autowrapper's scope, specifically impacting models like Kimi K2 that utilize such constructs.
• Improved Debuggability and Reproducibility: Enhanced the autowrapper's output consistency by switching from hash-based naming to an incrementing counter for wrapped functions and ensuring consistent sorting of arguments and return values in the generated code. This makes debugging and tracing much more predictable.
• Test Suite Refinement: Refactored existing autowrapper tests for improved clarity and maintainability. A new test case, test_function_variadic, was added to specifically validate the autowrapper's handling of functions with variadic arguments (*args and **kwargs).

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[gemini-code-assist]

Code Review

This pull request fixes an important bug related to recognizing all argument types in function definitions for autowrapping, which improves model compatibility. The change to use a counter for wrapped function names is a good step towards reproducibility. The test suite has been significantly refactored for clarity and robustness.

I've found a couple of minor issues:

1. The code for collecting function arguments can be made more concise.
2. The test refactoring revealed a bug in the code generation logic where wrapped functions can contain unreachable return statements.

I've left comments with suggestions for these points. Overall, this is a solid contribution.

[brian-dellabetta]

very impressive

[shanjiaz]

Looks good to me!