Fix nested def in @triton.jit to raise intended UnsupportedLanguageConstruct

[swjng]

New contributor declaration

• I am not making a trivial change, such as fixing a typo in a comment.

• I have written a PR description following these rules.

• I have run pre-commit run --from-ref origin/main --to-ref HEAD.

• Select one of the following.

  • I have added tests.
    • /python/test for end-to-end tests

• Select one of the following.

  • I have not added any lit tests.

---

Problem

Defining a nested function inside a @triton.jit kernel surfaces an opaque error:

@triton.jit
def kernel(ptr, n, BLOCK: tl.constexpr):
    def combine(a, b):   # <-- nested def
        return a + b
    ...

Before this change:

triton.compiler.errors.CompilationError: at 2:16:
def kernel(ptr, n, BLOCK: tl.constexpr):
    def combine(a, b):
                ^
StopIteration()

The message StopIteration() gives the user no indication of what went wrong.

Root cause

visit_FunctionDef in code_generator.py called self.visit(node.args) before checking if self.fn:. For a nested def, visit_arg calls

next(p for p in self.jit_fn.params if p.name == node.arg)

which raises StopIteration because the nested function's parameters are not in the outer JIT function's param list. The visitor's generic exception handler then wraps StopIteration as a CompilationError with the bare repr as its message.

Fix

Move the if self.fn: guard before self.visit(node.args) so that UnsupportedLanguageConstruct is raised immediately with a message directing the user to move helper functions to module level.

Nested def was already an unsupported construct (the guard existed, just in the wrong order). This change only improves the error message.

Test

test_err_nested_function_def in python/test/unit/language/test_compile_errors.py:

• asserts StopIteration does not appear in the error message
• asserts the error message mentions nested function