feat(autodiff): rewrite the broadcast-and-scale gradient region - #133
Open
AlekseiChirkovVention wants to merge 3 commits into
Open
feat(autodiff): rewrite the broadcast-and-scale gradient region#133AlekseiChirkovVention wants to merge 3 commits into
AlekseiChirkovVention wants to merge 3 commits into
Conversation
Pin the seven-clause predicate that identifies a broadcast-then-divide region by what it computes, and the five-node matmul region that replaces it. The predicate is semantic rather than origin-based: no artifact records which rule emitted a node, so a hand-built chain carrying no origin marker is asserted to expand exactly as a generated one does. The negative cases each assert the program is returned unchanged and unmutated with nothing raised, because broadcast and divide are general operators whose legitimate uses this pass must not disturb.
The helper seeded its set from the program's inherited type index, which the rewrite carries through unchanged. The broadcast intermediate therefore always survived in it, so the criterion that only that value is removed could never be observed, and the two preservation tests could not both hold. Seed from the values the nodes actually produce and read. No assertion changed.
A training step lowers a forward graph and the derivative program generated from it, so expanding only the forward mean leaves a backend without reduction, broadcast, and division handlers still unable to compile one step. Rewrite the matching broadcast-then-divide region into the same matmul-against-constants shape the forward pass emits. The region is identified by what it computes, not by where it came from. No artifact records which rule emitted a node, so a claim to recognise a generated chain would be unprovable; instead a seven-clause predicate decidable from the artifact alone selects any region for which the substitution is exact, and an algebraic identity holds it up for every value of the right shape. A chain failing any clause is returned untouched and raises nothing. Broadcast and divide are general operators with legitimate uses, so declining is the correct response to a near miss -- unlike an unsupported mean, which is unambiguously inside a pass named for mean expansion and fails closed instead. The exact-divisor clause is a scope guard rather than a correctness guard: the rewrite preserves value for any divisor, and the clause confines the pass to the reduction case and keeps the one inexact substitution on an exact integer element count.
code-tc
approved these changes
Aug 28, 2026
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
3 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.
Rewrite the broadcast-and-scale region of a derivative program into the same
matmuls-against-constants shape the forward pass emits, so one handler registry lowers a
forward graph and the derivative generated from it.
Goal
Expanding only the forward mean delivers nothing usable: a training step lowers both
artifacts, so a backend without reduction, broadcast, and division handlers still cannot
compile one step.
expand_mean_derivative_programcloses that gap.Why the region is identified by what it computes, not by where it came from
No artifact carries an origin marker — nodes do not record which rule emitted them — so any
claim to recognise "the generated mean gradient" would be unprovable, and naming it that way
would promise a guarantee the code cannot honour.
Instead a seven-clause predicate, decidable from the artifact alone, selects any region for
which the substitution is exact, and an algebraic identity holds it up: for any
gof shape[1, 1], broadcastinggto[r, c]and multiplying ones[r,1]bygand then by ones[1,c]produce the same value at every position, whatever producedg.The exact-divisor clause is a scope guard, not a correctness guard, and the module says so
rather than implying a safety it does not provide. The rewrite preserves value for any
divisor. The clause earns its place by confining the pass to the reduction case it is named
for, and by keeping the one inexact step — substituting
· (1/d)for/ d— on a divisorthat is an exact integer element count the framework itself generates, rather than an
arbitrary application-chosen one whose numerics a caller may depend on.
Declining is silent, and deliberately asymmetric
A chain failing any clause is returned untouched with nothing raised. Broadcast and divide are
general operators with many legitimate uses, so declining is the correct response to a near
miss. That is the opposite of this feature's behaviour on an unsupported mean, which fails
closed — a mean is unambiguously inside the declared domain of a pass named for mean
expansion, and silently leaving one would be surprising.
Test evidence
Level: unit. Tests first: the failing set landed in
9b19716;54f2ee8made it pass.27 cases, all passing. The focused autodiff suite moved 792 → 819, exactly the new items.
Broad non-integration suite: 1070 passed, 9 failed, 1 skipped — the 9 pre-existing and
unrelated, unchanged in count across this series.
Nine tests carry the positive burden; the other 18 are near-miss negatives. That asymmetry is
worth knowing when reading them: a predicate that matched nothing would also satisfy all 18,
so the positive cases — including a hand-built chain carrying no origin marker, which must
expand identically to a generated one — are what prove the matcher is semantic and actually
fires.
Review confirmed the rewrite fires on real generated programs across both
keepdimsforms,several shapes including non-square in both orientations, and both differentiable dtypes.
515ef3eamends a measurement helper in the test file. Two preservation tests were mutuallyunsatisfiable as written: one asserts the program's inherited type index is carried through
unchanged, while the other measured which values survive from that same index — so the
broadcast intermediate was on both sides by construction and its removal could never be
observed. The helper now measures the values the nodes actually produce and read. No assertion
changed; the amended criterion is still falsifiable.
Stack position
Stack position 4 of 8 — based on
feat/matmul-based-mean-expansion.Cannot merge into
mainbeforefeat/generated-constant-operator,chore/generic-lowering-reference-consumer, andfeat/matmul-based-mean-expansion.Known follow-ups, deliberately not fixed here
An ill-typed program violating the declared field types of its own dataclass can still surface
a bare builtin rather than a categorized error. Unreachable from any framework path, and the
same shape as a residual already parked on an earlier branch — worth fixing once, module-wide,
rather than in three places.