Autodiff coverage: close silent dispatch gaps + KSP-generated coverage guard#774
Merged
Conversation
Differentiability is a three-layer contract (@Diff -> generated backward contract -> trace dispatch). KSP locked the first two; the dispatch was hand-maintained and unguarded, so elu/leakyRelu/permute had correct backward formulas that were never wired -> their grads silently fell to null. Close the gaps and make the drift impossible to reintroduce: - Wire elu/leakyRelu/permute; fix permuteBackward's axes decode (trace records List<Int>, not IntArray). - Make cos/sin/tril/gather/indexSelect/unfold/convTranspose1d differentiable: @Diff + backward formula + finite-diff test each. Structural backwards are raw-loop adjoints (no new forward ops, no converters). sin/cos/convTranspose1d become abstract so the tracing wrapper records them. - Document sign/ge/lt/convert as non-differentiable by design. - KSP now emits DifferentiableTensorOpsRules.ruleNames; the dispatch when is refactored into a name-keyed backwardDispatch map exposing dispatchedOpNames; AutodiffCoverageTest asserts ruleNames is a subset of dispatchedOpNames. - operators.json carries isDifferentiable (+ optional diffRuleName), validated by validateOperatorSchema. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
This was referenced Jun 29, 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
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 #773.
Differentiability is a three-layer contract that must agree per op:
@Diff(contract) → a generatedDifferentiableTensorOps.<rule>Backward(formula) → a dispatch arm inbuildBackwardFromTrace(wiring). KSP already made a missing formula a compile error, but the wiring was hand-maintained in a downstream module and unguarded — soelu,leakyReluandpermutehad correct backward formulas that were never dispatched, and their gradients were silently dropped tonull.Close the gaps
elu/leakyRelu/permuteinto the dispatch. FixpermuteBackward's attribute decode (the tracer recordsaxesasList<Int>, notIntArray— it would have thrown the first time it ran, confirming it never had).cos/sin/tril/gather/indexSelect/unfold/convTranspose1ddifferentiable:@Diff+ a backward formula + a finite-difference parity test each. The structural ones are raw-loop adjoints inDefaultExecutionTape(likeconv1dBackward) — no new forward primitives, no StableHLO converter changes.sin/cos/convTranspose1dbecome abstract onTensorOpsso the tracing wrapper records them (both concrete backends already implement them).sign/ge/lt/convertas non-differentiable by design.Make the drift impossible (the core change)
TracingWrapperProcessornow also emitsDifferentiableTensorOpsRules.ruleNames— the authoritative set of@Diffrule names.whenis refactored into a name-keyedbackwardDispatchmap exposingdispatchedOpNames.AutodiffCoverageTestassertsruleNames ⊆ dispatchedOpNames— failing CI the instant any@Diffop lacks a wired backward (it would have caught the elu/leakyRelu/permute bug).OperatorDocProcessornow emitsisDifferentiable(+ optionaldiffRuleName) per function inoperators.json, with the schema extended andvalidateOperatorSchemakeeping it well-formed — the manifest becomes the single source of differentiability truth.Tests
AutodiffCoverageTest(dispatch covers every@Diffop).OpsAutodiffBackwardTest— finite-diff parity for all 10 newly-correct ops.jvmTestgreen forskainet-lang-core,skainet-compile-dag,skainet-compile-core,skainet-backend-cpu,skainet-compile-hlo;validateOperatorSchemagreen.Considered & rejected
Generating the dispatch table itself from
@Diff(zero hand-wiring): blocked by@DiffSOURCEretention + the tape living in a downstream module (KSP inskainet-lang-corecan't see those annotations from, or emit into,skainet-compile-dag). The generated rule-set + coverage test gives the same no-drift guarantee today.🤖 Generated with Claude Code