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

CHANGELOG.md

@@ -4,11 +4,19 @@
 
 ### Fixed
 
+- **A same-dtype `PrecisionBoundary` leaves no cast in a trace** (`skainet-lang-core`, SKEEP-008 phase A2, #1433): `stage("body", dtype = FP32)` inside an FP32 network called `convert`, which returns its input on the CPU backend but is still recorded by the tracing wrapper, so the exported module carried two `stablehlo.convert` `f32 -> f32`. The boundary now passes a tensor of the target dtype through without calling `convert`.
 - **`concat` backward with inputs of different sizes** (`skainet-compile-dag`, #1490): the rule 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 got another input's slice, or `null`, without an error. Each input now receives its own `narrow` of the upstream at the offset the forward gave it, for any sizes and a negative `dim`.
 - **`squeeze` backward with no axis** (`skainet-compile-dag`, #1490): `squeeze(x)` removes every size-1 axis, but the backward unsqueezed dim 0 alone, so `x: [1, 3, 1]` got a `[1, 3]` gradient. The gradient is now the upstream reshaped to the input's shape.
 
+### Removed
+
+- **`dense<TLayer>()`, `stage<TStage>()` and the mixed-precision scaffold** (`skainet-lang-core`, SKEEP-008 phase A4, #1433): the generic overloads of `NeuralNetworkDsl`, deprecated at `ERROR` level in 0.63.0, are gone; use `dense(outputDimension, dtype = FP16) { }` / `stage(id, dtype = FP16) { }`. `MixedPrecisionModule`, `InternalMixedPrecisionModule` and `MixedPrecisionTensorOps`, deprecated in 0.63.0 and never implemented, are removed, together with `MixedPrecisionException`, which only they could throw. A third-party `NeuralNetworkDsl` implementation drops its three overrides.
+
 ### Added
 
+- **Training through a mixed-precision boundary** (`skainet-lang-core`, `skainet-compile-dag`, SKEEP-008 phase A3, #1433): `TensorOps.convert` is a `@Diff` op. Its 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 and the layers behind `dense(dtype = …)` / `stage(id, dtype = …)` train; a cast from or to an integer dtype has no gradient. `MixedPrecisionBackwardTest` checks the rule, a finite-difference gradient through FP16 boundaries, parity with the all-FP32 network, and an SGD run. The inference-only notes are gone; the ops matrix Backward column reads 60/60.
+- **SKEEP-008 acceptance tests** (`skainet-lang-core`, `skainet-backend-cpu`, #1433): `repackTo(Dense(4))` is checked bit-exact for Q5_0, Q5_1, Q5_K, Q6_K and `BITNET_B1_58` too, and one FP32 network with both a BF16 `dense` and an FP16 `stage` runs against the FP32 reference.
+- **Mixed-precision export test** (`skainet-compile-hlo`, SKEEP-008 phase A2, #1433): `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 exported module.
 - **`scaledDotProductAttention` backward** (`skainet-compile-dag`, #1384): the attention adjoint (dQ, dK, dV and the dense mask's gradient) replaces the stub that returned `null`, so attention layers train. It honours `scale`, `causal`, sliding `window` / `rightContext`, `keyLengths` and the dense mask, recomputes the probabilities instead of storing them, and reduces grouped-query heads into their shared K/V. Masked scores get exactly zero gradient. `SdpaBackwardTest` checks every option against finite differences and three cases against PyTorch autograd; `stubbedBackwardOpNames` is empty and the coverage matrix's Backward column is complete.
 
 ## [0.63.0] - 2026-11-06

docs/modules/ROOT/pages/explanation/mixed-precision.adoc

@@ -102,8 +102,9 @@ ctx.ops.convert(inner.forward(x, ctx), outerType)  // cast back out
 Three properties follow from using `TensorOps.convert` — an ordinary tensor op — rather than
 hiding the conversion inside the layer:
 
-* *Free when degenerate.* `convert` returns its input unchanged when the dtypes already
-  match, so a boundary between equal dtypes costs nothing.
+* *Free when degenerate.* A tensor that already has the target dtype is passed through
+  without a `convert` call, so a boundary between equal dtypes costs nothing and leaves no
+  cast in a trace or an export.
 * *Visible in traces and exports.* The tracing wrappers record `convert` like any other op,
   and the StableHLO converter lowers it to `stablehlo.convert` — a compiled or exported graph
   shows exactly where precision changes, and a later optimisation pass can fold adjacent
@@ -115,14 +116,55 @@ hiding the conversion inside the layer:
 The dtype is passed as a *value* (`dense(256, dtype = BF16)`), not as a type argument. The
 earlier API sketched `dense<BF16>(256)`, but an interface method's type parameter cannot be
 `reified`, so the implementation had no `KClass` to build the layer's tensors with — those
-overloads could never be implemented and are now deprecated. Passing the `DType` object gives
+overloads could never be implemented and have been removed. Passing the `DType` object gives
 the implementation both the instance (for `convert`) and `dtype::class` (for tensor creation).
 
+== What an export looks like
+
+Tracing a network and exporting it to StableHLO gives one `stablehlo.convert` pair per
+boundary, with everything between them typed in the narrow format. For the network above
+without its `relu`, so that the BF16 layer is directly followed by the FP16 stage, the
+module has this shape (abridged: constants, transposes and bias adds are left out):
+
+[source,mlir]
+----
+%v6  = stablehlo.convert %arg0 : (tensor<1x784xf32>) -> tensor<1x784xbf16>   // into the BF16 dense
+%v10 = stablehlo.dot_general %v6, %v7, ... -> tensor<1x256xbf16>
+%v13 = stablehlo.convert %v12 : (tensor<1x256xbf16>) -> tensor<1x256xf32>    // and back out
+%v14 = stablehlo.convert %v13 : (tensor<1x256xf32>) -> tensor<1x256xf16>     // into the FP16 stage
+%v15 = stablehlo.dot_general %v14, %v8, ... -> tensor<1x64xf16>
+%v18 = stablehlo.dot_general %v17, %v9, ... -> tensor<1x10xf16>
+%v21 = stablehlo.convert %v20 : (tensor<1x10xf16>) -> tensor<1x10xf32>       // and back out
+----
+
+The `bf16 -> f32` cast followed directly by `f32 -> f16` is the price of two adjacent
+boundaries with different inner dtypes: the network's own dtype is always restored in
+between. `MixedPrecisionExportTest` pins this output.
+
+== Training through a boundary
+
+`convert` is differentiable between floating-point dtypes, so a mixed-precision layer or
+stage trains like any other. Its backward rule treats the cast as the identity and casts the
+gradient back to the source dtype:
+
+* where the gradient *leaves* a narrow region (the cast back out, seen from the loss side
+  first) it is rounded to the narrow format, exactly as the forward values were on the way in;
+* where it *enters* the outer network again it is widened, which is exact.
+
+The weights behind a boundary are stored in the narrow dtype, their gradients have that dtype
+too, and the optimizer updates them in place. That is a faithful simulation of training in the
+narrow format, with its limits: an update smaller than half a unit in the last place of a
+weight is lost. Keeping FP32 master weights and scaling the loss to avoid that is automatic
+mixed precision, which is out of scope (see below). A cast from or to an integer dtype has no
+gradient.
+
+One thing to know when checking gradients by hand: `input(n)` clears `requiresGrad` on the
+tensor it receives unless asked not to (`input(n, requiresGrad = true)`). With the flag
+cleared, the first cast has nothing to propagate and the network *input* gets no gradient; the
+parameters are not affected.
+
 == What it does not do (yet)
 
-* *Training.* `convert` is not differentiable today, so gradients stop at a boundary: a
-  mixed-precision layer or stage is inference-only. SKEEP-008 phase A3 makes `convert` a
-  `@Diff` op, which lifts this.
 * *Automatic mixed precision.* There is no autocast policy, no FP32 master weights, no loss
   scaling. The dtype is always an explicit, local choice in the code.
 * *Integer layers.* A cast to Int8 truncates; it is not quantization. Integer inference goes

docs/modules/ROOT/pages/reference/ops-status-matrix.adoc

@@ -72,7 +72,7 @@ Rows are `Operator.function` pairs. A gap is marked in source with `@NotImplemen
 | `TensorOps.lt` | ✖ | — | ❌ | — | — | ✅ | — 
 | `TensorOps.ge` | ✖ | — | ❌ | — | — | ✅ | — 
 | `TensorOps.tril` | ✖ | — | ❌ | ✅ | — | ✅ | — 
-| `TensorOps.convert` | ✖ | — | ✅ | — | — | ✅ | — 
+| `TensorOps.convert` | ✖ | — | ✅ | ✅ | — | ✅ | — 
 | `TensorOps.gather` | ✖ | — | ✅ | ✅ | — | ✅ | — 
 | `TensorOps.indexSelect` | ✖ | — | ✅ | ✅ | — | ✅ | — 
 | `TensorOps.exp` | ✖ | — | ✅ | ✅ | — | ✅ | — 
@@ -83,7 +83,7 @@ Rows are `Operator.function` pairs. A gap is marked in source with `@NotImplemen
 | `TensorOps.scaledDotProductAttention` | ✖ | composite · decompose | ✅ | ✅ | — | ✅ | — 
 | `Similarity.cosineDistance` | ✖ | — | ❌ | — | ✅ | ✅ | ✅ 
 
-| *Done* | *1 / 67* | *3 / 67* | *51 / 67* | *59 / 59* | *1 / 67* | *67 / 67* | *1 / 67* 
+| *Done* | *1 / 67* | *3 / 67* | *51 / 67* | *60 / 60* | *1 / 67* | *67 / 67* | *1 / 67* 
 |===
 
 Per-function detail including notes lives in xref:reference/operators/generated/index.adoc[Operator reference].

docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc

@@ -1,12 +1,28 @@
 = SKEEP-008: Precision and encoding conversion
 :description: SKaiNET proposal for the two ways a tensor's numeric form changes after it exists: a per-layer or per-stage dtype in the network DSL (mixed precision), and re-encoding stored weights between dense and packed formats (repackTo and an explicit, lossy requantizeTo).
 
-Status: Draft +
+Status: Implemented +
 Audience: SKaiNET maintainers and contributors; users of the network DSL; loader and kernel authors who consume `TensorEncoding` +
 Created: 2026-10-21 +
 Tracking issue: https://github.com/SKaiNET-developers/SKaiNET/issues/1433[#1433] +
 Origin: the stub inventory behind the operator coverage matrix (https://github.com/SKaiNET-developers/SKaiNET/pull/1370[#1370]). Six `TODO()` stubs in `NetworkBuilder.kt` (`dense<TLayer>()` ×2, `stage<TStage>()`, in both DSL implementations) and `TensorStorage.repackTo`, which throws for every encoding change, were left without a tracking issue because both change public API semantics.
 
+== Implementation status
+
+[cols="1,3,2",options="header"]
+|===
+| Phase | What shipped | Where
+
+| B1 | Lossless `repackTo`, lossy `requantizeTo` | https://github.com/SKaiNET-developers/SKaiNET/pull/1470[#1470], 0.63.0
+| A1 | `PrecisionBoundary`, `dense(dtype = …)`, `stage(id, dtype = …)`, deprecations | https://github.com/SKaiNET-developers/SKaiNET/pull/1469[#1469], 0.63.0
+| A2 | Trace and StableHLO export pinned by `MixedPrecisionExportTest`; a same-dtype boundary no longer records a cast | https://github.com/SKaiNET-developers/SKaiNET/pull/1492[#1492]
+| A3 | `convert` is `@Diff`; gradient check and training through a boundary in `MixedPrecisionBackwardTest` | https://github.com/SKaiNET-developers/SKaiNET/pull/1492[#1492]
+| A4 | Generic `dense<T>()` / `stage<T>()` overloads and the `MixedPrecision*` scaffold removed | https://github.com/SKaiNET-developers/SKaiNET/pull/1492[#1492]
+|===
+
+The sections below are the proposal as accepted. Where the implementation differs, the text
+says so; phase 1 restrictions that A3 lifted are kept for the record.
+
 == Summary
 
 SKaiNET has two public entry points that promise to change a tensor's numeric form and do
@@ -181,8 +197,9 @@ public class PrecisionBoundary<TOuter : DType, TInner : DType, V>(
 }
 ----
 
-`convert` already returns its input when the dtype matches, so a boundary between equal dtypes
-costs nothing. Two consecutive boundaries of the same inner dtype are a cast followed by its
+A tensor that already has the target dtype is passed through without calling `convert` (the
+tracing wrapper would record the call even though the backend returns its input), so a boundary
+between equal dtypes costs nothing and leaves no cast in a trace. Two consecutive boundaries of the same inner dtype are a cast followed by its
 inverse; a later optimisation pass may fold them, and nothing in this design depends on that.
 
 Weights of a `TLayer` dense are created in `TLayer` by the existing initialisers. User-supplied
@@ -328,13 +345,19 @@ Each phase is a DARC feature with lane sub-issues, linked from the tracking issu
 
 . Should a `reified` extension (`dense<FP16>(64)`) be offered as sugar over the value form? It
   reads well, but extensions named like members interact badly with overload resolution and IDE
-  completion.
+  completion. +
+  *Resolved: no.* Only the value form shipped.
 . Should a stage boundary that ends in the same dtype as the next stage elide the cast pair at
-  DSL time, or leave that to a compile pass?
+  DSL time, or leave that to a compile pass? +
+  *Open.* The DSL elides only the degenerate case (inner dtype equal to the network dtype). Two
+  adjacent boundaries still export a `narrow -> f32` cast followed by `f32 -> narrow`; folding
+  that pair is left to a compile pass.
 . Is `TensorStorage` the right owner of `requantizeTo`, or should it be a free function in
-  `sk.ainet.lang.tensor.data` next to `Requantizers`, keeping `TensorStorage` a descriptor?
+  `sk.ainet.lang.tensor.data` next to `Requantizers`, keeping `TensorStorage` a descriptor? +
+  *Resolved:* it is a member of `TensorStorage`, next to `repackTo`.
 . Does phase 2 need per-boundary `DTypePolicy` metadata, so that
-  `DTypeConstraintResolutionPass` can see and validate the boundaries in a compiled graph?
+  `DTypeConstraintResolutionPass` can see and validate the boundaries in a compiled graph? +
+  *Open.* A3 did not need it: the boundaries are ordinary `convert` nodes in the graph.
 
 == References
 

docs/modules/skeep/pages/index.adoc

@@ -140,6 +140,6 @@ Every proposal should include:
 | The browser lane — packed SIMD128 kernels as a sidecar module, quantize-at-load for SafeTensors, and the common-code fixes that keep wasm off the reference kernel
 
 | xref:skeep:008-precision-and-encoding-conversion.adoc[SKEEP-008]
-| Draft
+| Implemented
 | Precision and encoding conversion — per-layer and per-stage dtype in the network DSL, lossless `repackTo` and lossy `requantizeTo` for stored weights
 |===

skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt

@@ -121,6 +121,40 @@ class MixedPrecisionDslTest {
         assertTrue(maxDiff > 0f, "FP16 rounding must change the result, or the cast never happened")
     }
 
+    @Test
+    fun bf16DenseAndFp16StageInOneNetwork_runAndRound() {
+        // The SKEEP-008 acceptance network: both boundary kinds in one FP32 network.
+        val mixed = definition {
+            network<FP32, Float>(ctx) {
+                input(2)
+                dense(3, dtype = BF16, id = "layer1") {
+                    weights { fromArray(w1) }
+                    bias { fromArray(b1) }
+                }
+                stage("head", dtype = FP16) {
+                    dense(2, "layer2") {
+                        weights { fromArray(w2) }
+                        bias { fromArray(b2) }
+                    }
+                }
+            }
+        }
+
+        val reference = fp32Reference(ctx).forward(input(), ctx)
+        val result = mixed.forward(input(), ctx)
+
+        assertEquals(Shape(1, 2), result.shape)
+        assertEquals(FP32::class, result.dtype, "the last boundary must cast back to the network dtype")
+        assertEquals(2, mixed.modules.filterIsInstance<PrecisionBoundary<*, *, *>>().size)
+        var maxDiff = 0f
+        for (c in 0 until 2) {
+            val diff = abs(result.data[0, c] - reference.data[0, c])
+            assertTrue(diff <= 5e-2f, "mixed network drifted too far at $c: diff=$diff")
+            maxDiff = maxOf(maxDiff, diff)
+        }
+        assertTrue(maxDiff > 0f, "narrow rounding must change the result, or the casts never happened")
+    }
+
     @Test
     fun sameDtypeBoundary_isExact() {
         // convert returns its input when the dtype matches, so an FP32 stage inside an

skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api

@@ -106,6 +106,7 @@ public final class sk/ainet/lang/graph/DefaultGradientTape : sk/ainet/lang/graph
 	public fun conv3dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List;
 	public fun convTranspose1dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List;
 	public fun convTranspose2dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List;
+	public fun convertBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List;
 	public fun copy ()Lsk/ainet/tape/ExecutionTape;
 	public fun cosBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List;
 	public fun divScalarBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List;

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt

@@ -1156,6 +1156,27 @@ public class DefaultGradientTape(
         return listOf(upstream.ops.tril(upstream, k))
     }
 
+    override fun convertBackward(upstream: Tensor<DType, Any>, output: Tensor<DType, Any>, inputs: List<Tensor<DType, Any>>, attributes: Map<String, Any?>): List<Tensor<DType, Any>?> {
+        // d(convert)/dx is taken as the identity: the rounding of a narrowing cast is passed
+        // straight through. The gradient goes back in the source dtype, so a gradient entering
+        // an FP16 / BF16 region is itself rounded to that format, as the forward values are.
+        val input = inputs[0]
+        // A backend may return its input for a same-dtype cast. That call is still recorded,
+        // with one tensor as both input and output; a gradient for it would be added to itself.
+        if (session.refOf(input).id == session.refOf(output).id) return listOf(null)
+        val sourceType = floatDTypeOf(input) ?: return listOf(null)
+        if (floatDTypeOf(output) == null) return listOf(null)
+        return listOf(upstream.ops.convert(upstream, sourceType))
+    }
+
+    /** The dtype instance of a floating-point tensor that `convert` can target, else `null`. */
+    private fun floatDTypeOf(tensor: Tensor<DType, Any>): DType? = when (tensor.dtype) {
+        sk.ainet.lang.types.FP32::class -> sk.ainet.lang.types.FP32
+        sk.ainet.lang.types.FP16::class -> sk.ainet.lang.types.FP16
+        sk.ainet.lang.types.BF16::class -> sk.ainet.lang.types.BF16
+        else -> null
+    }
+
     override fun gatherBackward(upstream: Tensor<DType, Any>, output: Tensor<DType, Any>, inputs: List<Tensor<DType, Any>>, attributes: Map<String, Any?>): List<Tensor<DType, Any>?> {
         // gather(input[vocab,emb], indices, dim=0) -> backward scatter-adds upstream rows into the gathered rows.
         // Differentiable w.r.t. input only; indices gradient is null.
@@ -1475,6 +1496,7 @@ public class DefaultGradientTape(
             "sin" to ::sinBackward,
             "cos" to ::cosBackward,
             "tril" to ::trilBackward,
+            "convert" to ::convertBackward,
             "gather" to ::gatherBackward,
             "indexSelect" to ::indexSelectBackward,
             "unfold" to ::unfoldBackward,