Skip to content

Autodiff coverage: close silent dispatch gaps + KSP-generated coverage guard#774

Merged
michalharakal merged 1 commit into
developfrom
feat/autodiff-coverage
Jun 28, 2026
Merged

Autodiff coverage: close silent dispatch gaps + KSP-generated coverage guard#774
michalharakal merged 1 commit into
developfrom
feat/autodiff-coverage

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Closes #773.

Differentiability is a three-layer contract that must agree per op: @Diff (contract) → a generated DifferentiableTensorOps.<rule>Backward (formula) → a dispatch arm in buildBackwardFromTrace (wiring). KSP already made a missing formula a compile error, but the wiring was hand-maintained in a downstream module and unguarded — so elu, leakyRelu and permute had correct backward formulas that were never dispatched, and their gradients were silently dropped to null.

Close the gaps

  • Wire elu / leakyRelu / permute into the dispatch. Fix permuteBackward's attribute decode (the tracer records axes as List<Int>, not IntArray — it would have thrown the first time it ran, confirming it never had).
  • Make cos / sin / tril / gather / indexSelect / unfold / convTranspose1d differentiable: @Diff + a backward formula + a finite-difference parity test each. The structural ones are raw-loop adjoints in DefaultExecutionTape (like conv1dBackward) — no new forward primitives, no StableHLO converter changes. sin/cos/convTranspose1d become abstract on TensorOps so the tracing wrapper records them (both concrete backends already implement them).
  • Document sign / ge / lt / convert as non-differentiable by design.

Make the drift impossible (the core change)

  • TracingWrapperProcessor now also emits DifferentiableTensorOpsRules.ruleNames — the authoritative set of @Diff rule names.
  • The dispatch when is refactored into a name-keyed backwardDispatch map exposing dispatchedOpNames.
  • New AutodiffCoverageTest asserts ruleNames ⊆ dispatchedOpNames — failing CI the instant any @Diff op lacks a wired backward (it would have caught the elu/leakyRelu/permute bug).
  • OperatorDocProcessor now emits isDifferentiable (+ optional diffRuleName) per function in operators.json, with the schema extended and validateOperatorSchema keeping it well-formed — the manifest becomes the single source of differentiability truth.

Tests

  • AutodiffCoverageTest (dispatch covers every @Diff op).
  • OpsAutodiffBackwardTest — finite-diff parity for all 10 newly-correct ops.
  • Full jvmTest green for skainet-lang-core, skainet-compile-dag, skainet-compile-core, skainet-backend-cpu, skainet-compile-hlo; validateOperatorSchema green.

Considered & rejected

Generating the dispatch table itself from @Diff (zero hand-wiring): blocked by @Diff SOURCE retention + the tape living in a downstream module (KSP in skainet-lang-core can'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

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>
@github-actions

Copy link
Copy Markdown

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

  • Operator documentation: docs/modules/operators/_generated_/
  • JSON schema output: operators.json

Artifacts:

  • Download the documentation-preview-774 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.

@michalharakal michalharakal merged commit c785151 into develop Jun 28, 2026
9 checks passed
@michalharakal michalharakal deleted the feat/autodiff-coverage branch June 28, 2026 21:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Autodiff coverage: close silent dispatch gaps + KSP-generated coverage guard

1 participant