feat[next]: Add support for tuple comprehensions (not stacked) - #2487
feat[next]: Add support for tuple comprehensions (not stacked)#2487SF-N wants to merge 91 commits into
Conversation
…porting nesting (extracted from GridTools#2487)
| 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 |
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.
…upport # Conflicts: # src/gt4py/next/ffront/foast_passes/type_deduction.py # src/gt4py/next/ffront/foast_to_gtir.py # tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py
- Parametrize the new concat_where tuple tests with the 'static_domains' fixture, matching the rest of the file. - Update the unsupported-construct catalogue entry for generator expressions: they are now supported as the argument of 'tuple(...)'. - Restructure the new generator-expression diagnostics per the error-messages guide (smallest span, rationale as notes, hints) and pin them in test_diagnostic_messages.py. - Only mention generator expressions in the type-constructor error when the constructor is 'tuple'.
`NestedTuple` and its siblings in `gt4py.eve.extended_typing` were
defined as plain assignments holding a string forward reference to
themselves:
NestedTuple = tuple[Union[_T_co, "NestedTuple[_T_co]"], ...]
Subscripting such an alias does not substitute the type parameter inside
the forward reference, so `NestedTuple[Foo]` expands to `tuple[Foo |
ForwardRef('NestedTuple[_T_co]'), ...]` and `get_type_hints()` fails
with `NameError: name '_T_co' is not defined` when resolving it in the
namespace of the annotated object. These aliases could therefore not be
used as datamodel/node field annotations at all, which is why #2487 had
to fall back to `target: Any` instead of `NestedTuple[DataSymbol]` in
`foast.TupleComprehensionMapper`.
Changes:
- The `Nested*` / `MaybeNested*` aliases are now defined with the PEP
695 `type` statement, which substitutes type parameters properly across
the recursion.
- `SimpleTypeValidatorFactory` breaks the resulting cycle. Such an alias
is well founded, unlike the `type A = A` cycle `eval_type_alias()`
rejects: a single resolution step already yields an annotation which is
not an alias itself. What it does not yield is a *finite* one, so the
alias occurring inside its own definition is now handed a deferred
validator, filled in as soon as the definition has been processed.
`test_recursive_type_alias_is_not_supported` was renamed to
`test_cyclic_type_alias_is_not_supported`, since only degenerate cycles
remain unsupported.
Once this is merged, the `target: Any` workaround in #2487 can be
removed.
**Disclaimer**: This PR and its description were written largely with
the help of AI. Code was reviewed briefly.
Now that eve supports type aliases recursing through a container (GridTools#2814), annotate the comprehension target precisely as 'MaybeNestedInTuple[DataSymbol]' — the 'Maybe' variant since a bare-name target yields a lone symbol. Align 'func_to_foast.parse_target' and the 'type_deduction' helper signatures with the same spelling.
|
Continued in stacked PR #2833 |
Adds support for tuple comprehensions, e.g. for usage on tracers.