Skip to content

[compiler] add basic lambda support#1958

Open
hinriksnaer wants to merge 5 commits intopytorch:mainfrom
hinriksnaer:enable-lambda
Open

[compiler] add basic lambda support#1958
hinriksnaer wants to merge 5 commits intopytorch:mainfrom
hinriksnaer:enable-lambda

Conversation

@hinriksnaer
Copy link
Copy Markdown
Collaborator

Summary

  • Add visit_Lambda to TypePropagation and WalkDeviceAST so lambda expressions can be used inside @helion.kernel functions
  • Lambdas are compiled into real types.FunctionType objects via compile() + eval() with the kernel's module globals
  • Enables hl.reduce(lambda a, b: a + b, ...) instead of requiring a separately defined named function

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Apr 6, 2026
@hinriksnaer hinriksnaer requested review from fulvius31 and jansel April 6, 2026 20:43
Comment on lines +2354 to +2356
source = f"lambda {ast.unparse(node.args)}: {ast.unparse(node.body)}"
code = compile(source, filename=node._location.filename, mode="eval")
func = eval(code, self.func.fn.__globals__)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This approach doesn't seam ideal. While it will handle basic cases it doesn't handle closures and also uses tracing so some other operations won't work.

I'd be ok with some basic support here, but we should make sure to give good error messages if people go outside what we support.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants