Skip to content

feat(nn): finish SKEEP-008 — export test, differentiable convert, scaffold removal (phases A2, A3, A4) - #1492

Merged
michalharakal merged 2 commits into
developfrom
feature/1433-skeep-008-export-autograd-cleanup
Nov 10, 2026
Merged

michalharakal merged 2 commits into
developfrom
feature/1433-skeep-008-export-autograd-cleanup

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Finishes SKEEP-008 (precision and encoding conversion). B1 (#1470) and A1 (#1469) shipped in 0.63.0; this PR carries the three remaining phases and flips the proposal to Status: Implemented.

Tracking issue: #1433 — it can be closed when this merges.

A2 — export

  • MixedPrecisionExportTest traces an FP32 network with one BF16 dense and one FP16 stage and pins the result: one convert pair per boundary in the graph, inner layers typed in the narrow dtype, and stablehlo.convert f32 -> bf16, bf16 -> f32, f32 -> f16, f16 -> f32 in the module. DtypeForwardPropagationPass leaves the boundaries alone.
  • Fix: a same-dtype boundary (stage("body", dtype = FP32) in an FP32 network) still recorded convert — the CPU backend returns its input, but the tracing wrapper records the call — so the export carried two stablehlo.convert f32 -> f32. PrecisionBoundary now passes a tensor of the target dtype through without calling convert.

A3 — autograd

  • TensorOps.convert is @Diff. convertBackward treats the cast as the identity and casts the upstream gradient back to the source dtype, so a gradient entering an FP16 / BF16 region is rounded to that format, as the forward values are.
  • A cast from or to an integer dtype has no gradient. A recorded cast whose input and output are one tensor returns none, otherwise its gradient would be added to itself.
  • MixedPrecisionBackwardTest (7 tests): the rule itself, central finite differences through FP16 boundaries, parity with the all-FP32 network for a BF16 dense + FP16 stage, and an SGD run that lowers the loss and updates every weight behind a boundary.
  • The inference-only notes are removed from the KDoc and mixed-precision.adoc, which gains "What an export looks like" and "Training through a boundary". Ops matrix Backward column: 60 / 60.
  • The SKEEP asked for a phase-1 test pinning "gradients stop at a boundary". It was never written and is not added here, since this PR removes the restriction.

A4 — cleanup

  • Removed the generic dense<TLayer>() ×2 and stage<TStage>() from NeuralNetworkDsl and both implementations (deprecated at ERROR in 0.63.0).
  • Removed MixedPrecisionModule, InternalMixedPrecisionModule, MixedPrecisionTensorOps (deprecated in 0.63.0, never implemented).
  • Also removed MixedPrecisionException, which lived in the same file and which only the scaffold could throw. It was not marked deprecated in 0.63.0.
  • The DType compatibility / promotion tests from MixedPrecisionTensorOpsTest move to types/DTypeCompatibilityTest.

Acceptance criteria

  • repackTo(Dense(4)) is now checked bit-exact for Q5_0, Q5_1, Q5_K, Q6_K and BITNET_B1_58 as well (synthetic blocks; no quantizer produces them).
  • One FP32 network with both a BF16 dense and an FP16 stage runs against the FP32 reference.
  • SKEEP-008: Status: Implemented, an implementation-status table, and Open Questions 1 and 3 marked resolved (2 and 4 stay open).

Notes for review

  • API: convertBackward added to DifferentiableTensorOps / DefaultGradientTape; the removals above. .api files are regenerated.
  • ops-status-matrix.adoc was edited by hand (the convert row and the total), not regenerated with generateDocs.
  • README is untouched: its "inference only" and "SKEEP-008 draft" lines are in "What's New in 0.63.0" and describe that release.

Test plan

  • ./gradlew apiCheck checkUnmarkedStubs
  • jvmTest for skainet-lang-core, skainet-lang-models, skainet-compile-dag, skainet-compile-hlo, skainet-compile-opt, skainet-backend-cpu
  • linuxX64 compile of skainet-lang-core, skainet-compile-dag and the skainet-compile-hlo tests
  • Native and wasm test targets (CI; not run locally)

🤖 Generated with Claude Code

michalharakal and others added 2 commits November 10, 2026 22:23
…ffold removal (#1433)

Phases A2, A3 and A4 of SKEEP-008 (precision and encoding conversion).

A2, export: MixedPrecisionExportTest traces an FP32 network with one BF16
dense and one FP16 stage and pins one convert pair per boundary in the
graph and stablehlo.convert f32->bf16, bf16->f32, f32->f16, f16->f32 in
the module. A same-dtype boundary recorded a convert although the backend
returns its input, so the export carried two f32->f32 casts;
PrecisionBoundary now passes a tensor of the target dtype through.

A3, autograd: TensorOps.convert is @Diff. The backward rule treats the
cast as the identity and casts the upstream gradient back to the source
dtype, so a gradient entering an FP16/BF16 region is rounded to that
format. A cast from or to an integer dtype has no gradient, and a recorded
same-tensor cast returns none, or its gradient would be added to itself.
MixedPrecisionBackwardTest: the rule, central finite differences through
FP16 boundaries, parity with the all-FP32 network, and an SGD run that
updates every weight behind a boundary. The inference-only notes are
removed; the ops matrix Backward column reads 60/60.

A4, cleanup: the generic dense<TLayer>() / stage<TStage>() overloads
(deprecated at ERROR in 0.63.0) and MixedPrecisionModule,
InternalMixedPrecisionModule, MixedPrecisionTensorOps are removed, with
MixedPrecisionException, which only they could throw. The DType tests
from MixedPrecisionTensorOpsTest move to DTypeCompatibilityTest.

Acceptance: repackTo(Dense(4)) is checked bit-exact for Q5_0, Q5_1, Q5_K,
Q6_K and BITNET_B1_58, and one network with both boundary kinds runs
against the FP32 reference. SKEEP-008 moves to Status: Implemented.

Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.2 <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-1492 artifact to view the complete documentation locally.

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

1 similar comment
@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-1492 artifact to view the complete documentation locally.

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

@michalharakal
michalharakal merged commit 6ba5caf into develop Nov 10, 2026
22 checks passed
@michalharakal
michalharakal deleted the feature/1433-skeep-008-export-autograd-cleanup branch November 10, 2026 22:05
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.

2 participants