feat[next]: Add support for tuple comprehensions#2487
Open
SF-N wants to merge 17 commits into
Open
Conversation
tehrengruber
changed the title
| new_type = types[index] | ||
| case ts.VarArgType(element_type=element_type): | ||
| new_type = ( | ||
| element_type # TODO: we only temporarily allow any index for vararg types |
Contributor
There was a problem hiding this comment.
This is for direct access to tracers[0] * factor, tracers[1] * factor, which I personally think is an anti pattern. I left it here until we take a decision on this. We could also make it an optional feature. One of the disadvantages is that it is not possible to fully type check the field operator at definition time, since the tuple length is only known at call / compile time. The user will then get an error in unroll_map_tuple.
There was a problem hiding this comment.
Pull request overview
This PR adds frontend + lowering support for tuple comprehensions written as tuple(<genexpr>), enabling mapped operations over tuple-typed inputs (including variadic tuple annotations) and lowering them through a new iterator builtin (map_tuple) that is later unrolled into explicit tuple element operations.
Changes:
- Parse
tuple(<generator expression>)into a dedicated FOAST node and pretty-print it. - Type-deduce tuple comprehensions (incl. variadic tuple annotations via
VarArgType) and lower them to iterator IR using a newmap_tuplebuiltin. - Add an iterator transform to unroll
map_tuplecalls, plus integration/unit tests for supported and unsupported cases.
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py | Adds negative tests for invalid tuple-comprehension forms. |
| tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py | Adds execution tests covering fixed/variadic/nested tuple comprehensions. |
| tests/next_tests/integration_tests/cases.py | Adds temporary allocation/size handling for VarArgType (currently fixed to length 3). |
| src/gt4py/next/type_system/type_translation.py | Adds variadic/generic tuple hint handling and tuple constructor typing. |
| src/gt4py/next/type_system/type_specifications.py | Introduces VarArgType to represent variadic tuples. |
| src/gt4py/next/type_system/type_info.py | Extends concretization logic to account for VarArgType. |
| src/gt4py/next/iterator/type_system/type_synthesizer.py | Adds type synthesizer for new map_tuple builtin. |
| src/gt4py/next/iterator/transforms/unroll_map_tuple.py | New transform to unroll map_tuple into explicit tuple construction. |
| src/gt4py/next/iterator/transforms/pass_manager.py | Wires UnrollMapTuple into iterator transform pipelines. |
| src/gt4py/next/iterator/builtins.py | Registers map_tuple as an iterator builtin. |
| src/gt4py/next/ffront/past_passes/type_deduction.py | Relaxes out= typing check to compatible types. |
| src/gt4py/next/ffront/func_to_foast.py | Parses tuple(genexpr) into FOAST TupleComprehension. |
| src/gt4py/next/ffront/foast_to_past.py | Uses concretizability check for out return type validation. |
| src/gt4py/next/ffront/foast_to_gtir.py | Lowers tuple comprehensions to map_tuple(lambda)(iterable) calls, incl. unpacking targets. |
| src/gt4py/next/ffront/foast_pretty_printer.py | Adds pretty-print support for tuple comprehensions. |
| src/gt4py/next/ffront/foast_passes/type_deduction.py | Adds typing rules for TupleComprehension and vararg tuple indexing behavior. |
| src/gt4py/next/ffront/field_operator_ast.py | Adds FOAST node types for tuple comprehensions and mapper structure. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
476
to
+487
| def _verify_builtin_type_constructor(self, node: ast.Call) -> None: | ||
| if len(node.args) > 0: | ||
| arg = node.args[0] | ||
| if not ( | ||
| isinstance(arg, ast.Constant) | ||
| or (isinstance(arg, ast.UnaryOp) and isinstance(arg.operand, ast.Constant)) | ||
| ): | ||
| raise errors.DSLError( | ||
| self.get_location(node), | ||
| f"'{self._func_name(node)}()' only takes literal arguments.", | ||
| ) | ||
| assert isinstance(node.func, ast.Name) | ||
| (arg,) = node.args | ||
| if not ( | ||
| isinstance(arg, ast.Constant) | ||
| or (isinstance(arg, ast.UnaryOp) and isinstance(arg.operand, ast.Constant)) | ||
| or (node.func.id == "tuple" and isinstance(arg, ast.GeneratorExp)) | ||
| ): | ||
| raise errors.DSLError( | ||
| self.get_location(node), | ||
| f"'{self._func_name(node)}()' only takes literal arguments or a generator expression.", | ||
| ) |
| assert all(isinstance(elem, ts.DataType) for elem in tuple_types) | ||
| return ts.TupleType(types=tuple_types) | ||
| elif isinstance(args, tuple) and len(args) == 2 and args[1] is Ellipsis: | ||
| return ts.VarArgType(element_type=from_type_hint_same_ns(args[0])) |
Comment on lines
+432
to
+434
| new_type = ( | ||
| element_type # TODO: we only temporarily allow any index for vararg types | ||
| ) |
Comment on lines
+710
to
+720
| try: | ||
| type_ = element_type | ||
| for i in path: | ||
| if not isinstance(type_, ts.TupleType) or len(type_.types) <= i: | ||
| raise IndexError() | ||
| type_ = type_.types[i] | ||
| return self.visit(target_el, refine_type=type_, **inner_kwargs) | ||
| except IndexError: | ||
| raise errors.DSLError( | ||
| target_el.location, f"Cannot unpack non-iterable '{type_}' object." | ||
| ) from None |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds support for tuple comprehensions, e.g. for usage on tracers.