feat[next]: Embedded field scan #1365
Conversation
|
Tests covering all cases are in |
tehrengruber
changed the title
src/gt4py/next/embedded/operators.py
Outdated
| scan_range = embedded_context.closure_column_range.get() | ||
| assert self.axis == scan_range[0] | ||
|
|
||
| return _scan(self.fun, self.forward, self.init, scan_range, args, kwargs) |
There was a problem hiding this comment.
Why not place the function here?
| axis: common.Dimension | ||
|
|
||
| def __call__(self, *args: common.Field | core_defs.Scalar, **kwargs: common.Field | core_defs.Scalar) -> common.Field: # type: ignore[override] # we cannot properly type annotate relative to self.fun | ||
| scan_range = embedded_context.closure_column_range.get() |
There was a problem hiding this comment.
| scan_range = embedded_context.closure_column_range.get() | |
| scan_range: NamedRange = embedded_context.closure_column_range.get() |
src/gt4py/next/embedded/operators.py
Outdated
| for hpos in embedded_common.iterate_domain(non_scan_domain): | ||
| scan_loop(hpos) | ||
| if len(non_scan_domain) == 0: | ||
| # if we don't have any dimension orthogonal to scan_axis, we need to do one scan_loop | ||
| scan_loop(()) |
There was a problem hiding this comment.
| for hpos in embedded_common.iterate_domain(non_scan_domain): | |
| scan_loop(hpos) | |
| if len(non_scan_domain) == 0: | |
| # if we don't have any dimension orthogonal to scan_axis, we need to do one scan_loop | |
| scan_loop(()) | |
| if len(non_scan_domain) == 0: | |
| # if we don't have any dimension orthogonal to scan_axis, we need to do one scan_loop | |
| scan_loop(()) | |
| else: | |
| for hpos in embedded_common.iterate_domain(non_scan_domain): | |
| scan_loop(hpos) |
src/gt4py/next/utils.py
Outdated
| def flatten_nested_tuple(value: tuple[_S | tuple, ...]) -> tuple[_S, ...]: | ||
| if isinstance(value, tuple): | ||
| return sum((flatten_nested_tuple(v) for v in value), start=()) # type: ignore[arg-type] # cannot properly express nesting | ||
| else: | ||
| return (value,) |
There was a problem hiding this comment.
How about this?
| def flatten_nested_tuple(value: tuple[_S | tuple, ...]) -> tuple[_S, ...]: | |
| if isinstance(value, tuple): | |
| return sum((flatten_nested_tuple(v) for v in value), start=()) # type: ignore[arg-type] # cannot properly express nesting | |
| else: | |
| return (value,) | |
| def tree_leaves(value: S_ | tuple[_S | tuple, ...]) -> Iterable[_S, ...]: | |
| if isinstance(value, tuple): | |
| for el in value: | |
| yield from flatten_nested_tuple(el) | |
| else: | |
| yield value |
This also works on values (which would already be useful in one of your use-cases) and has a name that is aligned with tree_map. Only mentioning it here, but meant for the follow up PR.
There was a problem hiding this comment.
I believe your version is better because it is lazy, but the functionality is the same. Only in my version I am lying about the typing because otherwise you cannot pass a tuple because it immediately gets matched to _S and then you get a typing error. I opted for the version with a typing error inside but now obviously you cannot pass a value without getting a typing error. The best version uses a protocol NestedSequence that Enrique found in some codebase. I'll add an issue about this cleanup.
| k_size, | ||
| ) | ||
| ) # i_size bigger than in the other argument | ||
| inp2_np = np.fromfunction(lambda i, k: k, shape=(i_size, k_size), dtype=float) |
There was a problem hiding this comment.
A small gibe: Is this an example where local-view is easier to comprehend?
Adds the scalar scan operator for embedded field view.