feat[next]: Implement embedded as_offset#2612
Conversation
# Conflicts: # tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py
There was a problem hiding this comment.
Pull request overview
This PR implements embedded as_offset support for NdArrayField execution, enabling dynamic offset field premaps in the embedded backend and un-xfailing related embedded tests.
Changes:
- Adds an embedded
as_offsetbuiltin implementation that builds a gather connectivity from integer offset fields. - Adds type-deduction validation for unsupported non-single-target offsets.
- Adds unit coverage for embedded
as_offsetpremapping and removes the dynamic-offset embedded xfail.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/gt4py/next/embedded/nd_array_field.py |
Registers and implements embedded as_offset connectivity construction. |
src/gt4py/next/ffront/foast_passes/type_deduction.py |
Adds compile-time rejection for unsupported as_offset offset shapes. |
tests/next_tests/definitions.py |
Enables dynamic-offset tests for embedded backends. |
tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py |
Adds embedded runtime tests for as_offset premap behavior and unsupported neighbor offsets. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
havogt
left a comment
havogt
left a comment
There was a problem hiding this comment.
Self-review notes (via /code-review skill, xhigh effort)
9 finder angles → 7 verifiers → sweep. Two candidates were refuted before posting:
- Missing
source ∈ offset_field.dimscheck —Domain.__getitem__already raises a cleanKeyError("No Dimension of type 'X' is present in the Domain.")(common.py:528-529), so a defensive check would be redundant. - Cartesian-only check is a backward-incompatible tightening — non-Cartesian
as_offsetnever compiled on any backend (gtfn/dace/roundtrip all hit downstream asserts; ICON4py only usesKoff). The new check just improves diagnostics.
Inline findings below. One finding does not fit any changed line (_hyperslice reachability) — recorded here:
Pre-existing: inverse_image raises bare numpy ValueError when offset range is fully disjoint
src/gt4py/next/embedded/nd_array_field.py:704 (in _hyperslice, unchanged code). When offset_field.domain[source_dim] is fully disjoint from the data field's range (e.g. f on I=[0,10), offset_field on I=[100,110)), select_mask is all-False, xp.nonzero returns empty, and xp.min(empty) raises ValueError: zero-size array to reduction operation minimum which has no identity. Pre-existing bug in _hyperslice; this PR's new as_offset builtin makes it more reachable from user code.
| domain = offset_field.domain | ||
| xp = offset_field.array_ns | ||
| src_range = domain[source_dim].unit_range | ||
| shape = tuple(len(src_range) if dim == source_dim else 1 for dim in domain.dims) |
There was a problem hiding this comment.
Same value-only Dimension.__eq__ issue as the predicate (see comment on fbuiltins.py:514): if domain.dims ever contained two dims sharing .value with different .kind, both axes would get len(src_range) instead of just the source axis.
Currently shielded by Domain's uniqueness check (common.py:495) which uses the same value-only equality — so unreachable today, but coupled to that invariant. If uniqueness is ever relaxed (e.g. to allow same-name dims of different kinds), this shape computation breaks silently.
There was a problem hiding this comment.
Won't fix in this PR. Shielded by Domain's uniqueness invariant (which uses the same value-only equality), so unreachable today. The predicate-level tightening in fbuiltins.is_cartesian_offset blocks the construction one layer up. Worth revisiting only if Domain uniqueness is ever relaxed.
|
|
||
| def _as_offset(offset: fbuiltins.FieldOffset, offset_field: NdArrayField) -> common.Connectivity: | ||
| if not fbuiltins.is_cartesian_offset(offset): | ||
| target_dims = ", ".join(d.value for d in offset.target) |
There was a problem hiding this comment.
The full error-message string (lines 899–903) is verbatim-duplicated with type_deduction.py:894–899. Unified messaging was intentional, but the contract is currently enforced by hand; a small fbuiltins._not_cartesian_message(offset) -> str shared by both raisers would enforce it by construction. Each call site still controls its own exception type and location wrapping. Minor — flag only if you expect future drift.
There was a problem hiding this comment.
Won't fix in this PR — the hand-unified messages are intentional. Easy follow-up if reviewers later prefer a shared helper, but the contract is only two raisers right now.
|
Re: the pre-existing |
No description provided.