feat[next]: Infer as_fieldop type without domain #1853
Conversation
This file contains 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
…ot yet working for tuples
SF-N
commented
Feb 13, 2025
| @@ -492,10 +492,7 @@ def visit_SetAt(self, node: itir.SetAt, *, ctx) -> None: | |||
| # probably just change the behaviour of the lowering. Until then we do this more | |||
| # complicated comparison. | |||
| if isinstance(target_type, ts.FieldType) and isinstance(expr_type, ts.FieldType): | |||
| assert ( | |||
| set(expr_type.dims).issubset(set(target_type.dims)) | |||
There was a problem hiding this comment.
I think the subset requirement is not necessary anymore, e.g. in the case
itir.SetAt(
domain=unstructured_domain,
expr=im.as_fieldop(
im.lambda_("it")(im.reduce("plus", 0.0)(im.deref("it"))),
unstructured_domain,
)(im.ref("inp")),
target=im.ref("out"),
)
with
im.sym("inp", float_vertex_v2e_k_field) and im.sym("out", float_vertex_k_field)
the expr_type is
[Dimension(value='Vertex', kind=<DimensionKind.HORIZONTAL: 'horizontal'>), Dimension(value='KDim', kind=<DimensionKind.VERTICAL: 'vertical'>), Dimension(value='V2E', kind=<DimensionKind.LOCAL: 'local'>)]
and the target_type is
[Dimension(value='Vertex', kind=<DimensionKind.HORIZONTAL: 'horizontal'>), Dimension(value='KDim', kind=<DimensionKind.VERTICAL: 'vertical'>)]
which is correct in my opinion.
cf. test_fencil_with_nb_field_input
SF-N
commented
Feb 13, 2025
| )(im.ref("inp1", float_i_field), im.ref("inp2", float_i_field)), | ||
| ts.TupleType(types=[float_i_field, float_i_field]), | ||
| )(im.ref("inp1", float_i_field), im.ref("inp2", float_j_field)), | ||
| ts.TupleType(types=[float_ij_field, float_ij_field]), |
There was a problem hiding this comment.
Is ts.TupleType(types=[float_ij_field, float_ij_field] correct here?
SF-N
commented
Feb 14, 2025
src/gt4py/next/common.py
Outdated
| return ( | ||
| sorted(unique_dims, key=lambda dim: (kind_order[dim.kind], dim.value)) | ||
| if unique_dims | ||
| else [] |
Comment on lines
386
to
398
| if any(shifts_results[i]): | ||
| for input_dim in input_dims: | ||
| for shift_tuple in shifts_results[ | ||
| i | ||
| ]: # Use shift tuple corresponding to the input field | ||
| final_dim = ( | ||
| resolve_shift(input_dim, shift_tuple) or input_dim | ||
| ) # If ther are no shifts, take input_dim | ||
| if final_dim not in seen: | ||
| seen.add(final_dim) | ||
| output_dims.append(final_dim) | ||
| else: | ||
| output_dims.extend(input_dims) |
There was a problem hiding this comment.
- No special casing for tuple of empty tuple
((),) - Do not add input dims for empty tuple (the argument is unused in that case). Add test case
| return final_target | ||
|
|
||
| if any(shifts_results[i]): | ||
| for input_dim in input_dims: |
There was a problem hiding this comment.
Change ordering to:
for shift in ...:
for input_dim in ...:
...
for readability.
SF-N
commented
Feb 25, 2025
Comment on lines
+415
to
+417
| position_dims=float_vertex_k_field.dims, | ||
| defined_dims=float_vertex_field.dims, | ||
| element_type=ts.ListType(element_type=ts.ScalarType(kind=ts.ScalarKind.FLOAT64)), |
There was a problem hiding this comment.
Why is the IteratorType like this?
Please carefully check this test as I am not 100% sure.
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.
This extends the GTIR type inference to infer the type of as_fieldop calls without a domain.
TODOs: