Repository navigation
fix(compile-dag): concat and squeeze backward on shapes they got wrong (#1490) - #1491
Merged
Merged
Conversation
#1490) concatBackward split the upstream gradient by one input's size and picked chunk `offset / size`, which only lines up when every input has the same extent along the concat axis; other shapes received another input's slice, or null from a caught exception, without an error. Each input now gets its own narrow of the upstream at the offset the forward placed it, for any sizes and a negative dim. squeezeBackward unsqueezed dim 0 alone, so squeeze(x) with no axis on x: [1, 3, 1] produced a [1, 3] gradient. Squeeze only drops size-1 axes, so the gradient is the upstream reshaped to the input's shape. ConcatSqueezeBackwardTest: exact gradients under a position-weighted loss for unequal concat (dim 0, dim 1, negative dim, three inputs), equal concat, squeeze with and without an axis, and a squeeze-then- concat chain. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
aharakal
approved these changes
Nov 10, 2026
Merged
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
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.
Closes #1490.
Two
@Diffops with a wired dispatch arm whose backward returned a wrong result without an error, so neither the coverage matrix norAutodiffCoverageTestcould see them.concatBackwardThe rule split the upstream by one input's size along
dimand picked chunkoffset / size. That only lines up when every input has the same extent along the axis. Forconcat([a: [2,1], b: [2,2]], dim = 1)under a position-weighted loss,bgot[1, 2, 4, 5]instead of[2, 3, 5, 6]; shapes wheresplitthrew were caught and turned intonull.Now each input gets
narrow(upstream, dim, offset, size)at the offset the forward placed it, offsets accumulated in input order, negativedimnormalised, and a size-0 input still getsnull.squeezeBackwardsqueeze(x)with no axis removes every size-1 axis, but the rule wasunsqueeze(upstream, dim ?: 0), sox: [1, 3, 1]got a[1, 3]gradient. Squeeze only drops size-1 axes, so the gradient is simply the upstream reshaped to the input's shape, whichever axes were removed.Tests
ConcatSqueezeBackwardTest(7 cases): exact gradients under a loss whose weight at each output position is its row-major index + 1, so every position's gradient is known. Unequal concat on dim 1, on dim 0 with three inputs, with a negative dim; equal concat as a regression;squeeze()andsqueeze(dim)shapes; a squeeze-then-concat chain.:skainet-compile:skainet-compile-dag:jvmTest124/124,apiCheckclean.Out of scope
converthas no backward, so gradients throughPrecisionBoundaryand the FP16/BF16 branches of layer/group norm are zero. That is SKEEP-008 phase A3 (#1433).🤖 Generated with Claude Code