llmcompressor.pipelines.sequential.ast_utils.auto_wrapper

def auto_wrap(self, tree: ast.Module) -> ast.Module:
    """
    Modify ast by automatically wrapping any untraceable code segments. Segments to
    wrap are determined through analysis of the code and basic pattern matching

    :param tree: module containing a definition to an original forward function
    :return: module with added wrapper function definitions and function calls
    """
    tree = self.visit(tree)
    for fn_def in self._wrapper_fn_defs:
        tree.body.insert(0, fn_def)

    return ast.fix_missing_locations(tree)

def visit_Call(self, node: ast.Call) -> ast.Call:
    """
    Wrap any functions which use (4) variadic arguments or (2) match the ignore list
    """
    # check for variadic starred
    if any(isinstance(elem, ast.Starred) for elem in node.args):
        return self._wrap_if_possible(node)

    # attempt to evaluate caller and check against ignore list
    try:
        caller = self._eval_expr(node.func)

    except Exception:
        caller = None

    finally:
        if (
            isinstance(caller, (FunctionType, MethodType))
            and caller.__name__ in self.ignore
        ):
            return self._wrap_if_possible(node)

    return super().generic_visit(node)

def visit_Delete(self, node: ast.Delete):
    """
    Remove any deleted names from `self._local_names`,
    which are used to determine function wrapper arguments
    """
    ret = super().generic_visit(node)

    for target in node.targets:
        if isinstance(target, ast.Name):
            self._local_names.remove(target.id)

    return ret

def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
    """
    Remove decorators which prevent forward function recompilation
    For example, add_start_docstrings_to_model_forward

    Because `_wrapper_fn_defs` are appended after `visit` finishes, this function
    will not affect wrapper functions

    :param node: function definition whose decorators will be stripped
    :return: function definition without decorators
    """
    node.decorator_list = [
        decorator_name
        for decorator_name in node.decorator_list
        if isinstance(decorator_name, ast.Name)
        and decorator_name.id in ("can_return_tuple",)  # modifies func signature
    ]

    if node.name == "forward":
        for arg in node.args.args:
            self._local_names.add(arg.arg)
        for arg in node.args.posonlyargs:
            self._local_names.add(arg.arg)
        for arg in node.args.kwonlyargs:
            self._local_names.add(arg.arg)
        if node.args.vararg:
            self._local_names.add(node.args.vararg.arg)
        if node.args.kwarg:
            self._local_names.add(node.args.kwarg.arg)
    return super().generic_visit(node)

def visit_If(self, node: ast.If) -> ast.If | ast.Assign:
    """
    Attempt to statically evaluate the condition of the `if` statement. If the
    condition can not be statically evaluated (1), then attmept to wrap the `if`
    statement

    :param node: `if` statement which may be wrapped
    :return: if the `if` statement cannot be statically evaluated, return the
        `if` statement with the condition replaced by `True` or `False`.
        Otherwise, return a wrapper function call
    """
    try:
        value = bool(self._eval_expr(node.test))

        # force a wrap if any assignments occur within the if statement
        for expr in ast.walk(node):
            if isinstance(expr, ast.NamedExpr):
                raise Exception("If statement contains assignment")

    except Exception:
        return self._wrap_if_possible(node)

    else:
        node.test = ast.Constant(value=value)
        return super().generic_visit(node)

def visit_Name(self, node: ast.Name):
    """
    Add any new names in `self._local_names`,
    which are used to determine function wrapper arguments
    """
    if isinstance(node.ctx, ast.Store):
        self._local_names.add(node.id)

    return super().generic_visit(node)

def visit_Tuple(self, node: ast.Tuple) -> ast.Tuple | ast.Call:
    """
    (3) Wrap any tuples which use starred unpacking
    """
    if any(isinstance(elem, ast.Starred) for elem in node.elts):
        return self._wrap_if_possible(node)

    return super().generic_visit(node)