api: Fix staggered sum indices - #3012
Open
mloubout wants to merge 3 commits into
Open
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #3012 +/- ##
==========================================
+ Coverage 83.68% 83.70% +0.02%
==========================================
Files 257 257
Lines 54711 54777 +66
Branches 4686 4689 +3
==========================================
+ Hits 45785 45853 +68
+ Misses 8115 8112 -3
- Partials 811 812 +1
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| Where `expr` sits along `dim`, or None if its terms disagree. | ||
| """ | ||
| try: | ||
| indices = {i.indices_ref[dim] for i in expr.args} if expr.is_Add \ |
Contributor
There was a problem hiding this comment.
ultra-nitpicking, I'm not fond of this style when u have to break to new line like that anyway
| class DimensionTuple(EnrichedTuple): | ||
|
|
||
| def __getitem_hook__(self, dim): | ||
| # An exact hit wins over a `_defines` overlap. A derived Dimension |
Contributor
There was a problem hiding this comment.
these AI-written comments are uselessly verbose and thus only confusing, I think the first sentence until the full stop is enough
| return self.getters[d] | ||
| raise KeyError | ||
|
|
||
| def dindex(self, dim): |
Contributor
There was a problem hiding this comment.
I've just realised that maybe you can likely call into __getitem__hook__ and then apply the .index later?
basically the implementation of this method and __getitem__hook__ are currently virtually the same
mloubout
force-pushed
the
fix-staggered-sum-indices
branch
from
August 28, 2026 16:01
e96d4ca to
6993249
Compare
`__getitem_hook__` matched on `_defines` overlap alone. A derived Dimension carries its parent in `_defines`, so for a Bundle indexed by `(p_rec, rp_recx)` -- `rp_recx` being a `CustomDimension` whose parent is `p_rec` -- the lookup for `rp_recx` matched the `p_rec` entry first and returned the number of sparse points where the number of interpolation weights was meant. That size becomes the innermost stride in `_generate_fsz`, so the receiver kernels of a vectorized Operator read `w[p*npoint + rp]` instead of `w[p*2 + rp]` and run off the end of the array. Observed as an out-of-bounds `__global__` read under compute-sanitizer and a run-to-run varying, sometimes NaN, elastic TTI gradient on CUDA. Try an exact hit before falling back to the overlap, in both `__getitem_hook__` and `dindex`.
mloubout
force-pushed
the
fix-staggered-sum-indices
branch
from
August 28, 2026 16:12
6993249 to
5580d46
Compare
`Add` reports its first argument's `indices_ref`, so a sum whose terms sit at different staggered locations names a position only one of them has, and `x0` gets resolved against it for all of them. The shear strain `v_x.dy + v_y.dx` of a staggered velocity is the canonical case: both terms land on the cell corner, so a shift onto it should be a no-op, and instead each picked up a spurious one. Differentiation is linear at every order, so split such a sum in `Derivative._eval_fd`. Relative error on `D(a+b)` against `D(a) + D(b)` was 0.63 at order 0, 1.20 at order 1 and 0.95 at order 2, with `expand=False` at order 2 returning exactly zero. `generic_derivative` also short-circuited a zeroth order derivative only when `x0` was empty, building a stencil around an expression already sitting at `x0`. `index_at` answers where an expression sits, and both call sites use it.
mloubout
force-pushed
the
fix-staggered-sum-indices
branch
from
August 28, 2026 17:11
5580d46 to
8e805ab
Compare
A LocalObject carries expressions in its constructor arguments and in its initializer, and both end up in the generated code, but FindApplications only visited Expressions, Iterations and Calls. Any macro they apply was therefore left undefined -- ROUND_UP, say, for an auto-padded stride reaching a plan descriptor.
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
2 participants
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.
Misc FD (and bundle) bug fixes I ran into with some adjoint/gradient tests