From 4feb9a720730bb2cfeb4a5ed564dc46133ce3363 Mon Sep 17 00:00:00 2001
From: michalharakal <michalharakal@users.noreply.github.com>
Date: Sun, 1 Nov 2026 11:34:02 +0100
Subject: [PATCH 1/3] feat(nn): mixed-precision dense/stage via
 PrecisionBoundary (SKEEP-008 A1, #1433)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Implement rollout phase A1 of SKEEP-008: the DSL takes the layer dtype as a
value (dense(64, dtype = BF16) { }, stage(id, dtype = FP16) { }) and wraps the
layer or block in PrecisionBoundary, which casts in and out with the existing
TensorOps.convert — free when dtypes match, visible in traces and StableHLO
exports, inference-only until convert becomes @Diff (phase A3).

The old generic overloads (dense<TLayer>(), stage<TStage>()) could never be
implemented — an interface method's type parameter cannot be reified, so there
was no KClass to build tensors with. They are deprecated at level ERROR and
their six TODO() stubs replaced by error() bodies; the NetworkBuilder.kt entry
leaves config/stub-allowlist.txt. The never-implemented scaffold
(MixedPrecisionModule, InternalMixedPrecisionModule, MixedPrecisionTensorOps)
is deprecated for removal.

Docs: new Antora explanation article on the value and limits of mixed
precision, linked from the nav and from SKEEP-008.
---
 config/stub-allowlist.txt                     |   1 -
 docs/modules/ROOT/nav.adoc                    |   1 +
 .../pages/explanation/mixed-precision.adoc    | 141 ++++++++++
 .../lang/nn/dsl/MixedPrecisionDslTest.kt      | 168 ++++++++++++
 .../api/jvm/skainet-lang-core.api             |  14 +
 .../sk/ainet/lang/nn/MixedPrecisionModule.kt  |   8 +
 .../sk/ainet/lang/nn/PrecisionBoundary.kt     |  37 +++
 .../sk/ainet/lang/nn/dsl/NetworkBuilder.kt    | 241 ++++++++++++++----
 .../tensor/ops/MixedPrecisionTensorOps.kt     |   4 +
 9 files changed, 570 insertions(+), 45 deletions(-)
 create mode 100644 docs/modules/ROOT/pages/explanation/mixed-precision.adoc
 create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt
 create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt

diff --git a/config/stub-allowlist.txt b/config/stub-allowlist.txt
index 6c5f75509..36fe986e3 100644
--- a/config/stub-allowlist.txt
+++ b/config/stub-allowlist.txt
@@ -27,5 +27,4 @@ skainet-io/skainet-io-iree-params/src/wasmJsMain/kotlin/sk/ainet/io/irpa/FileBac
 skainet-io/skainet-io-iree-params/src/wasmWasiMain/kotlin/sk/ainet/io/irpa/FileBackedTransfer.wasmwasi.kt|1|IRPA FileBacked transfer on wasmWasi (follow-up of #1423)
 
 # --- Lang core: DSL and tensor views ---------------------------------------------------------
-skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt|6|SKEEP-008 (#1433) mixed-precision dense<T>() / stage<T>()
 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/dsl/TensorDSL.kt|1|#1417 TensorDefineDslImpl.create() (unused; implement or delete)
diff --git a/docs/modules/ROOT/nav.adoc b/docs/modules/ROOT/nav.adoc
index 4ce2f0986..aabdb2bef 100644
--- a/docs/modules/ROOT/nav.adoc
+++ b/docs/modules/ROOT/nav.adoc
@@ -44,6 +44,7 @@
 ** xref:explanation/kernel-selection.adoc[Kernel SPI and the selection algorithm]
 ** xref:explanation/schedules.adoc[Algorithm and schedule]
 ** xref:explanation/quantization-process.adoc[The quantization process]
+** xref:explanation/mixed-precision.adoc[Mixed precision]
 ** xref:explanation/theory/index.adoc[Mathematical theory]
 *** xref:explanation/theory/matmul.adoc[Matrix multiplication]
 ** xref:explanation/examples/index.adoc[Worked examples]
diff --git a/docs/modules/ROOT/pages/explanation/mixed-precision.adoc b/docs/modules/ROOT/pages/explanation/mixed-precision.adoc
new file mode 100644
index 000000000..0ce5203ac
--- /dev/null
+++ b/docs/modules/ROOT/pages/explanation/mixed-precision.adoc
@@ -0,0 +1,141 @@
+= Mixed Precision: Why a Layer's Dtype Is Worth Choosing
+:description: Why SKaiNET lets one layer or one stage run in a narrower dtype than the rest of the network, what that buys (and doesn't buy yet), and how the PrecisionBoundary design keeps every cast visible.
+
+A network built with the {framework_name} DSL has one dtype: `sequential<FP32, Float> { … }`
+fixes it for every layer. Mixed precision (SKEEP-008, tracked in
+https://github.com/SKaiNET-developers/SKaiNET/issues/1433[#1433]) loosens that — one layer or
+one block can run in another dtype while the network around it stays unchanged:
+
+[source,kotlin]
+----
+val model = sequential<FP32, Float> {
+    input(784)
+    dense(256, dtype = BF16) { }     // one BF16 layer in an FP32 network
+    activation { it.relu() }
+    stage("head", dtype = FP16) {    // a whole FP16 block
+        dense(64) { }
+        dense(10)
+    }
+}
+----
+
+This page explains why that is worth having, what it costs, and where the current
+implementation stops.
+
+== Precision is not uniformly needed
+
+The observation behind every mixed-precision system — PyTorch AMP, JAX's dtype promotion,
+ONNX Runtime's float16 transforms — is that a network's tolerance for rounding is wildly
+uneven. The matmul-heavy bulk of a model (dense and convolution layers) tolerates 16-bit
+floats well: their error statistics average out over thousands of multiply-accumulates. A few
+places do not: softmax and normalization (small differences of large numbers), accumulators,
+and the loss. Forcing the whole network to the precision its most sensitive operation needs
+wastes memory and bandwidth on the 95% that doesn't care.
+
+Per-layer `dense(dtype = …)` and per-block `stage(id, dtype = …)` are exactly that dial: keep
+the skeleton FP32, drop the bulk into a cheaper dtype, and keep the choice *local* — a
+narrow layer is opt-in per call, never a network-wide default.
+
+== What the two 16-bit floats trade away
+
+[cols="1,2,2",options="header"]
+|===
+| | *FP16* | *BF16*
+
+| Layout
+| 1 sign, 5 exponent, 10 mantissa bits
+| 1 sign, 8 exponent, 7 mantissa bits
+
+| Range
+| overflows past 65 504
+| FP32's full range
+
+| Precision
+| ~3 decimal digits
+| ~2 decimal digits
+
+| Failure mode
+| overflow to `Inf` in activations with large magnitudes
+| coarse rounding; sums of many small values drift
+|===
+
+BF16 keeps FP32's exponent, so anything that fits in FP32 fits in BF16 — it trades mantissa
+for safety and is usually the first thing to try. FP16 is more precise per bit but needs the
+values to stay small. The `convert` codecs handle overflow (to infinity) and NaN correctly,
+but a model can still *silently* degrade — which is why the dtype appears in the code at the
+call site, not in a config file.
+
+== Why it matters for SKaiNET's targets
+
+{framework_name} runs on Kotlin Multiplatform targets where the usual mixed-precision
+arguments are amplified:
+
+* *Memory-constrained hosts.* A browser tab (wasm), an Android app, an iOS app — environments
+  where halving or quartering a model's working set decides whether it runs at all.
+* *Bandwidth-bound inference.* Most inference time is spent moving weights, not multiplying
+  them. Fewer bytes per element means fewer bytes through the cache hierarchy per token.
+* *Models that arrive mixed.* GGUF and ONNX models in the wild already mix precisions —
+  quantized matmul weights next to FP32 norms. A DSL that can only express one dtype has no
+  faithful home for them.
+
+[IMPORTANT]
+====
+Today the memory and bandwidth benefits are *not yet realized*: FP16 and BF16 tensors are
+float-backed (the values are rounded to the 16-bit format, the storage is still 4 bytes per
+element). Real 2-byte storage belongs to the storage model (SKEEP-003). What mixed precision
+gives you now is the *numerics* — a faithful simulation of the narrow format, the DSL surface,
+and exports that carry the casts — so models can be validated against 16-bit behaviour before
+the storage catches up, and the API doesn't change when it does.
+====
+
+== The design: one visible boundary
+
+The DSL wraps a mixed-precision layer or stage in a single module, `PrecisionBoundary`. Its
+forward pass is two casts around the wrapped module:
+
+[source,kotlin]
+----
+val x = ctx.ops.convert(input, innerType)    // cast in
+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.
+* *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
+  cast pairs.
+* *Honest about cost.* Every boundary copies its input and its output. A network that
+  switches dtype every layer can be *slower* than all-FP32, especially while narrow dtypes
+  are float-backed. Few, coarse boundaries (a `stage` around a block) beat many fine ones.
+
+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
+the implementation both the instance (for `convert`) and `dtype::class` (for tensor creation).
+
+== 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
+  through packed encodings and the quantized kernels — see
+  xref:explanation/quantization-process.adoc[The quantization process]. Mixed precision and
+  quantization are complementary, not competing: one changes the *float format* a layer
+  computes in, the other changes the *encoding* weights are stored and multiplied in.
+
+== Where to look
+
+* `sk.ainet.lang.nn.PrecisionBoundary` — the boundary module.
+* `NeuralNetworkDsl.dense(outputDimension, dtype, …)` and `stage(id, dtype, …)` — the DSL
+  entry points.
+* `TensorOps.convert` — the cast, including FP16/BF16 rounding through the format codecs.
+* xref:skeep:008-precision-and-encoding-conversion.adoc[SKEEP-008] — the full design,
+  including the storage-side `repackTo`/`requantizeTo` split and the phased rollout.
diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt
new file mode 100644
index 000000000..8a3e34dc6
--- /dev/null
+++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt
@@ -0,0 +1,168 @@
+package sk.ainet.lang.nn.dsl
+
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.context.ExecutionContext
+import sk.ainet.context.data
+import sk.ainet.lang.nn.Module
+import sk.ainet.lang.nn.PrecisionBoundary
+import sk.ainet.lang.nn.definition
+import sk.ainet.lang.nn.network
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.tensor.dsl.tensor
+import sk.ainet.lang.types.BF16
+import sk.ainet.lang.types.FP16
+import sk.ainet.lang.types.FP32
+import kotlin.math.abs
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/**
+ * SKEEP-008 phase A1 (#1433): `dense(dtype = …)` and `stage(id, dtype = …)` run their layers
+ * in another dtype behind a [PrecisionBoundary], casting in and out with `TensorOps.convert`.
+ */
+class MixedPrecisionDslTest {
+
+    private val ctx = DirectCpuExecutionContext()
+
+    // Values deliberately not representable in FP16/BF16, so the boundary's rounding is observable.
+    private val w1 = floatArrayOf(
+        0.1234567f, -1.7654321f,
+        2.3456789f, 0.0123456f,
+        -0.9876543f, 1.1111111f,
+    )
+    private val b1 = floatArrayOf(0.1111111f, -0.2222222f, 0.3333333f)
+    private val w2 = floatArrayOf(
+        1.9999999f, -0.1111111f, 0.5555555f,
+        -1.3333333f, 0.7777777f, -0.4444444f,
+    )
+    private val b2 = floatArrayOf(0.0101010f, -0.0202020f)
+
+    private fun input(): Tensor<FP32, Float> = data<FP32, Float>(ctx) {
+        tensor<FP32, Float> {
+            shape(1, 2) { fromArray(floatArrayOf(0.3f, -1.1f)) }
+        }
+    }
+
+    private fun fp32Reference(context: ExecutionContext): Module<FP32, Float> = definition {
+        network<FP32, Float>(context) {
+            input(2)
+            dense(3, "layer1") {
+                weights { fromArray(w1) }
+                bias { fromArray(b1) }
+            }
+            dense(2, "layer2") {
+                weights { fromArray(w2) }
+                bias { fromArray(b2) }
+            }
+        }
+    }
+
+    @Test
+    fun denseWithDtype_runsAndRoundsThroughBf16() {
+        val mixed = definition {
+            network<FP32, Float>(ctx) {
+                input(2)
+                dense(3, dtype = BF16, id = "layer1") {
+                    weights { fromArray(w1) }
+                    bias { fromArray(b1) }
+                }
+                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 boundary must cast back to the network dtype")
+        var maxDiff = 0f
+        for (c in 0 until 2) {
+            val diff = abs(result.data[0, c] - reference.data[0, c])
+            assertTrue(diff <= 2e-2f, "BF16 layer drifted too far at $c: diff=$diff")
+            maxDiff = maxOf(maxDiff, diff)
+        }
+        assertTrue(maxDiff > 0f, "BF16 rounding must change the result, or the cast never happened")
+    }
+
+    @Test
+    fun stageWithDtype_runsAndRoundsThroughFp16() {
+        val mixed = definition {
+            network<FP32, Float>(ctx) {
+                input(2)
+                stage("body", dtype = FP16) {
+                    dense(3, "layer1") {
+                        weights { fromArray(w1) }
+                        bias { fromArray(b1) }
+                    }
+                    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 boundary must cast back to the network dtype")
+        var maxDiff = 0f
+        for (c in 0 until 2) {
+            val diff = abs(result.data[0, c] - reference.data[0, c])
+            assertTrue(diff <= 2e-2f, "FP16 stage drifted too far at $c: diff=$diff")
+            maxDiff = maxOf(maxDiff, diff)
+        }
+        assertTrue(maxDiff > 0f, "FP16 rounding must change the result, or the cast never happened")
+    }
+
+    @Test
+    fun sameDtypeBoundary_isExact() {
+        // convert returns its input when the dtype matches, so an FP32 stage inside an
+        // FP32 network reproduces the reference bit-exactly.
+        val wrapped = definition {
+            network<FP32, Float>(ctx) {
+                input(2)
+                stage("body", dtype = FP32) {
+                    dense(3, "layer1") {
+                        weights { fromArray(w1) }
+                        bias { fromArray(b1) }
+                    }
+                    dense(2, "layer2") {
+                        weights { fromArray(w2) }
+                        bias { fromArray(b2) }
+                    }
+                }
+            }
+        }
+
+        val reference = fp32Reference(ctx).forward(input(), ctx)
+        val result = wrapped.forward(input(), ctx)
+
+        for (c in 0 until 2) {
+            assertEquals(reference.data[0, c], result.data[0, c], "FP32 boundary must be a no-op at $c")
+        }
+    }
+
+    @Test
+    fun mixedLayers_appearAsPrecisionBoundaryModules() {
+        val mixed = definition {
+            network<FP32, Float>(ctx) {
+                input(2)
+                dense(3, dtype = BF16, id = "layer1") {
+                    weights { fromArray(w1) }
+                    bias { fromArray(b1) }
+                }
+            }
+        }
+
+        val boundaries = mixed.modules.filterIsInstance<PrecisionBoundary<*, *, *>>()
+        assertEquals(1, boundaries.size, "the BF16 dense must be wrapped in one PrecisionBoundary")
+        assertEquals("layer1", boundaries.single().name)
+    }
+}
diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api
index 89f1aabab..e5b20c4e6 100644
--- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api
+++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api
@@ -2997,6 +2997,12 @@ public final class sk/ainet/lang/nn/NeuralNetworkExecutionContextKt {
 	public static final fun definition (Lkotlin/jvm/functions/Function2;)Lsk/ainet/lang/nn/Module;
 }
 
+public final class sk/ainet/lang/nn/PrecisionBoundary : sk/ainet/lang/nn/Module {
+	public fun <init> (Ljava/lang/String;Lsk/ainet/lang/nn/Module;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)V
+	public fun getModules ()Ljava/util/List;
+	public fun getName ()Ljava/lang/String;
+}
+
 public final class sk/ainet/lang/nn/TrainingRunnerKt {
 	public static final fun trainStep (Lsk/ainet/lang/nn/Module;Lsk/ainet/lang/nn/loss/Loss;Lsk/ainet/lang/nn/optim/Optimizer;Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
 }
@@ -3526,10 +3532,12 @@ public abstract interface class sk/ainet/lang/nn/dsl/NeuralNetworkDsl : sk/ainet
 	public static synthetic fun conv3d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IZLjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
 	public abstract fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module;
 	public abstract fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)V
+	public abstract fun dense (ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
 	public abstract fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module;
 	public abstract fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
 	public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module;
 	public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
+	public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
 	public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module;
 	public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
 	public abstract fun flatten (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
@@ -3553,6 +3561,7 @@ public abstract interface class sk/ainet/lang/nn/dsl/NeuralNetworkDsl : sk/ainet
 	public static synthetic fun softmax$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;ILjava/lang/Object;)V
 	public abstract fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module;
 	public abstract fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
+	public abstract fun stage (Ljava/lang/String;Lsk/ainet/lang/types/DType;Lkotlin/jvm/functions/Function1;)V
 	public abstract fun upsample2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
 	public abstract fun upsample2d (Lkotlin/Pair;Lsk/ainet/lang/tensor/ops/UpsampleMode;ZLjava/lang/String;)V
 	public static synthetic fun upsample2d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
@@ -3569,6 +3578,7 @@ public final class sk/ainet/lang/nn/dsl/NeuralNetworkDsl$DefaultImpls {
 	public static synthetic fun conv3d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IZLjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
 	public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module;
 	public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
+	public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
 	public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module;
 	public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
 	public static synthetic fun flatten$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
@@ -3597,6 +3607,7 @@ public final class sk/ainet/lang/nn/dsl/NeuralNetworkDslImpl : sk/ainet/lang/nn/
 	public final fun create ()Lsk/ainet/lang/nn/Module;
 	public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module;
 	public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)V
+	public fun dense (ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
 	public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module;
 	public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
 	public fun flatten (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
@@ -3618,6 +3629,7 @@ public final class sk/ainet/lang/nn/dsl/NeuralNetworkDslImpl : sk/ainet/lang/nn/
 	public fun softmax (ILjava/lang/String;)V
 	public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module;
 	public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
+	public fun stage (Ljava/lang/String;Lsk/ainet/lang/types/DType;Lkotlin/jvm/functions/Function1;)V
 	public fun upsample2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
 	public fun upsample2d (Lkotlin/Pair;Lsk/ainet/lang/tensor/ops/UpsampleMode;ZLjava/lang/String;)V
 }
@@ -3638,6 +3650,7 @@ public final class sk/ainet/lang/nn/dsl/StageImpl : sk/ainet/lang/nn/dsl/NeuralN
 	public final fun create ()Lsk/ainet/lang/nn/Module;
 	public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module;
 	public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)V
+	public fun dense (ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
 	public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module;
 	public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
 	public fun flatten (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
@@ -3661,6 +3674,7 @@ public final class sk/ainet/lang/nn/dsl/StageImpl : sk/ainet/lang/nn/dsl/NeuralN
 	public fun softmax (ILjava/lang/String;)V
 	public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module;
 	public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
+	public fun stage (Ljava/lang/String;Lsk/ainet/lang/types/DType;Lkotlin/jvm/functions/Function1;)V
 	public fun upsample2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
 	public fun upsample2d (Lkotlin/Pair;Lsk/ainet/lang/tensor/ops/UpsampleMode;ZLjava/lang/String;)V
 }
diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt
index 818993cab..4e7afd561 100644
--- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt
+++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt
@@ -39,6 +39,10 @@ import sk.ainet.lang.types.DType
  * }
  * ```
  */
+@Deprecated(
+    message = "Never implemented; replaced by PrecisionBoundary (SKEEP-008, #1433). Will be removed.",
+    replaceWith = ReplaceWith("PrecisionBoundary", "sk.ainet.lang.nn.PrecisionBoundary"),
+)
 public abstract class MixedPrecisionModule<TInput : DType, TOutput : DType, V>(
     protected val inputType: TInput,
     protected val outputType: TOutput,
@@ -210,6 +214,10 @@ public abstract class MixedPrecisionModule<TInput : DType, TOutput : DType, V>(
  * @param TInternal The internal computation precision type
  * @param V The value type corresponding to the DType
  */
+@Deprecated(
+    message = "Never implemented; replaced by PrecisionBoundary (SKEEP-008, #1433). Will be removed.",
+    replaceWith = ReplaceWith("PrecisionBoundary", "sk.ainet.lang.nn.PrecisionBoundary"),
+)
 public abstract class InternalMixedPrecisionModule<TIO : DType, TInternal : DType, V>(
     private val ioType: TIO,
     private val internalType: TInternal,
diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt
new file mode 100644
index 000000000..fa83e6cca
--- /dev/null
+++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt
@@ -0,0 +1,37 @@
+package sk.ainet.lang.nn
+
+import sk.ainet.context.ExecutionContext
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.types.DType
+
+/**
+ * Runs [inner], a module built in another dtype, inside a network of dtype [TOuter]:
+ * the input is cast to [innerType], [inner] runs, and the result is cast back to
+ * [outerType] (SKEEP-008, #1433).
+ *
+ * The casts use [sk.ainet.lang.tensor.ops.TensorOps.convert], which returns its input
+ * unchanged when the dtypes already match, so a boundary between equal dtypes costs
+ * nothing. `convert` is recorded by the tracing wrappers and lowers to
+ * `stablehlo.convert`, so the boundary stays visible in traces and exports.
+ *
+ * `convert` is not differentiable (SKEEP-008 phase A3 lifts this), so gradients stop at
+ * the boundary: the wrapped layers do not train.
+ */
+public class PrecisionBoundary<TOuter : DType, TInner : DType, V>(
+    override val name: String,
+    private val inner: Module<TInner, V>,
+    private val innerType: TInner,
+    private val outerType: TOuter,
+) : Module<TOuter, V>() {
+
+    @Suppress("UNCHECKED_CAST")
+    override val modules: List<Module<TOuter, V>>
+        // The dtype parameter is erased at runtime; the inner module is listed so that
+        // reflection, parameter collection and checkpointing still see its subtree.
+        get() = listOf(inner as Module<TOuter, V>)
+
+    override fun onForward(input: Tensor<TOuter, V>, ctx: ExecutionContext): Tensor<TOuter, V> {
+        val x = ctx.ops.convert(input, innerType)
+        return ctx.ops.convert(inner.forward(x, ctx), outerType)
+    }
+}
diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt
index bab5b93ef..fb910c4d0 100644
--- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt
+++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt
@@ -11,6 +11,7 @@ import sk.ainet.lang.nn.Input
 import sk.ainet.lang.nn.Linear
 import sk.ainet.lang.nn.MaxPool2d
 import sk.ainet.lang.nn.Module
+import sk.ainet.lang.nn.PrecisionBoundary
 import sk.ainet.lang.nn.Upsample2d
 import sk.ainet.lang.nn.normalization.BatchNormalization
 import sk.ainet.lang.nn.normalization.GroupNormalization
@@ -166,27 +167,39 @@ public interface NeuralNetworkDsl<T : DType, V> : NetworkDslItem {
     public fun gru(hiddenSize: Int, id: String = "", content: GRU<T, V>.() -> Unit = {})
 
     /**
-     * Creates a dense layer with precision override and specified output dimension.
-     * This allows individual layers to use different precision than the network default.
+     * Creates a dense layer whose weights and computation use [dtype] while the network stays
+     * in its own dtype: the layer's input is cast to [dtype] and its output cast back
+     * (SKEEP-008, #1433). The cast is `TensorOps.convert`, which is not differentiable yet,
+     * so a mixed-precision layer is inference-only: gradients stop at its boundary.
      *
-     * @param TLayer The precision type for this specific layer
      * @param outputDimension The number of neurons/output features
+     * @param dtype The dtype the layer's weights and computation use (e.g. `FP16`, `BF16`)
      * @param id Optional identifier for the layer
-     * @param content Configuration block for weights, bias, and activation
+     * @param content Configuration block for weights, bias, and activation, typed in [dtype]
      */
+    public fun <TLayer : DType> dense(
+        outputDimension: Int,
+        dtype: TLayer,
+        id: String = "",
+        content: DENSE<TLayer, V>.() -> Unit = {},
+    )
+
+    @Deprecated(
+        message = "The dtype type argument cannot be implemented (no KClass at runtime, " +
+            "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }",
+        level = DeprecationLevel.ERROR,
+    )
     public fun <TLayer : DType> dense(
         outputDimension: Int,
         id: String = "",
         content: DENSE<TLayer, V>.() -> Unit = {}
     ): Module<T, V>
 
-    /**
-     * Creates a dense layer with precision override without specifying output dimension.
-     *
-     * @param TLayer The precision type for this specific layer
-     * @param id Optional identifier for the layer
-     * @param content Configuration block where output dimension, weights, and bias are set
-     */
+    @Deprecated(
+        message = "The dtype type argument cannot be implemented (no KClass at runtime, " +
+            "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }",
+        level = DeprecationLevel.ERROR,
+    )
     public fun <TLayer : DType> dense(
         id: String = "",
         content: DENSE<TLayer, V>.() -> Unit = {}
@@ -438,18 +451,26 @@ public interface NeuralNetworkDsl<T : DType, V> : NetworkDslItem {
     public fun stage(id: String, content: NeuralNetworkDsl<T, V>.() -> Unit)
 
     /**
-     * Creates a precision-scoped stage within the network.
-     * This allows grouping layers with a specific precision type that differs
-     * from the network default, enabling fine-grained mixed-precision control.
-     *
-     * Unlike the basic stage method, this version allows changing the precision
-     * type for all layers within the stage scope.
+     * Creates a stage whose layers all use [dtype] while the network stays in its own dtype:
+     * the stage's input is cast to [dtype] and its output cast back (SKEEP-008, #1433).
+     * The cast is `TensorOps.convert`, which is not differentiable yet, so a mixed-precision
+     * stage is inference-only: gradients stop at its boundary.
      *
-     * @param TStage The precision type for all layers within this stage
      * @param id Identifier for the stage
-     * @param content DSL block containing layers with TStage precision
-     * @return Module that handles the precision conversion automatically
+     * @param dtype The dtype all layers in this stage use (e.g. `FP16`, `BF16`)
+     * @param content DSL block containing the stage's layers, typed in [dtype]
      */
+    public fun <TStage : DType> stage(
+        id: String,
+        dtype: TStage,
+        content: NeuralNetworkDsl<TStage, V>.() -> Unit,
+    )
+
+    @Deprecated(
+        message = "The dtype type argument cannot be implemented (no KClass at runtime, " +
+            "SKEEP-008); pass the dtype as a value: stage(id, dtype = FP16) { ... }",
+        level = DeprecationLevel.ERROR,
+    )
     public fun <TStage : DType> stage(
         id: String,
         content: NeuralNetworkDsl<TStage, V>.() -> Unit
@@ -1321,28 +1342,90 @@ public class StageImpl<T : DType, V>(
 
     override fun <TLayer : DType> dense(
         outputDimension: Int,
+        dtype: TLayer,
         id: String,
-        content: DENSE<TLayer, V>.() -> Unit
-    ): Module<T, V> {
-        // Create a mixed-precision module that handles conversion
-        TODO("Mixed-precision dense implementation needed")
+        content: DENSE<TLayer, V>.() -> Unit,
+    ) {
+        val inputDimension = lastDimension
+        lastDimension = outputDimension
+        currentShape = intArrayOf(outputDimension)
+        val layerName = getDefaultName(id, "linear", modules.size)
+        val impl = DenseImpl<TLayer, V>(
+            executionContext,
+            inputDimension = inputDimension,
+            _outputDimension = outputDimension,
+            id = layerName,
+            kClass = dtype.witnessOf(),
+        )
+        impl.content()
+        val innerModules = impl.create()
+        val inner = if (innerModules.size == 1) innerModules[0] else MLP(*innerModules.toTypedArray(), name = layerName)
+        modules += PrecisionBoundary(
+            name = layerName,
+            inner = inner,
+            innerType = dtype,
+            outerType = dtypeInstanceOf(kClass),
+        )
     }
 
+    override fun <TStage : DType> stage(
+        id: String,
+        dtype: TStage,
+        content: NeuralNetworkDsl<TStage, V>.() -> Unit,
+    ) {
+        val stageImpl = StageImpl<TStage, V>(executionContext, id, dtype.witnessOf())
+        stageImpl.lastDimension = lastDimension
+        stageImpl.currentShape = currentShape?.copyOf()
+        stageImpl.content()
+        lastDimension = stageImpl.lastDimension
+        currentShape = stageImpl.currentShape?.copyOf()
+        modules += PrecisionBoundary(
+            name = id,
+            inner = stageImpl.create(),
+            innerType = dtype,
+            outerType = dtypeInstanceOf(kClass),
+        )
+    }
+
+    @Deprecated(
+        message = "The dtype type argument cannot be implemented (no KClass at runtime, " +
+            "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }",
+        level = DeprecationLevel.ERROR,
+    )
     override fun <TLayer : DType> dense(
+        outputDimension: Int,
         id: String,
         content: DENSE<TLayer, V>.() -> Unit
-    ): Module<T, V> {
-        // Create a mixed-precision module that handles conversion
-        TODO("Mixed-precision dense implementation needed")
-    }
+    ): Module<T, V> = error(
+        "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " +
+            "dense(outputDimension, dtype = FP16, id) { ... }"
+    )
 
+    @Deprecated(
+        message = "The dtype type argument cannot be implemented (no KClass at runtime, " +
+            "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }",
+        level = DeprecationLevel.ERROR,
+    )
+    override fun <TLayer : DType> dense(
+        id: String,
+        content: DENSE<TLayer, V>.() -> Unit
+    ): Module<T, V> = error(
+        "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " +
+            "dense(outputDimension, dtype = FP16, id) { ... }"
+    )
+
+    @Deprecated(
+        message = "The dtype type argument cannot be implemented (no KClass at runtime, " +
+            "SKEEP-008); pass the dtype as a value: stage(id, dtype = FP16) { ... }",
+        level = DeprecationLevel.ERROR,
+    )
     override fun <TStage : DType> stage(
         id: String,
         content: NeuralNetworkDsl<TStage, V>.() -> Unit
-    ): Module<T, V> {
-        // Create a mixed-precision stage that handles conversion
-        TODO("Mixed-precision stage implementation needed")
-    }
+    ): Module<T, V> = error(
+        "Mixed-precision stage was redesigned (SKEEP-008): pass the dtype as a value, " +
+            "stage(id, dtype = FP16) { ... }"
+    )
 
     override fun batchNorm(
         numFeatures: Int,
@@ -1715,28 +1798,90 @@ public class NeuralNetworkDslImpl<T : DType, V>(
 
     override fun <TLayer : DType> dense(
         outputDimension: Int,
+        dtype: TLayer,
         id: String,
-        content: DENSE<TLayer, V>.() -> Unit
-    ): Module<T, V> {
-        // Create a mixed-precision module that handles conversion
-        TODO("Mixed-precision dense implementation needed")
+        content: DENSE<TLayer, V>.() -> Unit,
+    ) {
+        val inputDimension = lastDimension
+        lastDimension = outputDimension
+        currentShape = intArrayOf(outputDimension)
+        val layerName = getDefaultName(id, "linear", modules.size)
+        val impl = DenseImpl<TLayer, V>(
+            executionContext,
+            inputDimension = inputDimension,
+            _outputDimension = outputDimension,
+            id = layerName,
+            kClass = dtype.witnessOf(),
+        )
+        impl.content()
+        val innerModules = impl.create()
+        val inner = if (innerModules.size == 1) innerModules[0] else MLP(*innerModules.toTypedArray(), name = layerName)
+        modules += PrecisionBoundary(
+            name = layerName,
+            inner = inner,
+            innerType = dtype,
+            outerType = dtypeInstanceOf(kClass),
+        )
+    }
+
+    override fun <TStage : DType> stage(
+        id: String,
+        dtype: TStage,
+        content: NeuralNetworkDsl<TStage, V>.() -> Unit,
+    ) {
+        val stageImpl = StageImpl<TStage, V>(executionContext, id, dtype.witnessOf())
+        stageImpl.lastDimension = lastDimension
+        stageImpl.currentShape = currentShape?.copyOf()
+        stageImpl.content()
+        lastDimension = stageImpl.lastDimension
+        currentShape = stageImpl.currentShape?.copyOf()
+        modules += PrecisionBoundary(
+            name = id,
+            inner = stageImpl.create(),
+            innerType = dtype,
+            outerType = dtypeInstanceOf(kClass),
+        )
     }
 
+    @Deprecated(
+        message = "The dtype type argument cannot be implemented (no KClass at runtime, " +
+            "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }",
+        level = DeprecationLevel.ERROR,
+    )
     override fun <TLayer : DType> dense(
+        outputDimension: Int,
         id: String,
         content: DENSE<TLayer, V>.() -> Unit
-    ): Module<T, V> {
-        // Create a mixed-precision module that handles conversion
-        TODO("Mixed-precision dense implementation needed")
-    }
+    ): Module<T, V> = error(
+        "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " +
+            "dense(outputDimension, dtype = FP16, id) { ... }"
+    )
 
+    @Deprecated(
+        message = "The dtype type argument cannot be implemented (no KClass at runtime, " +
+            "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }",
+        level = DeprecationLevel.ERROR,
+    )
+    override fun <TLayer : DType> dense(
+        id: String,
+        content: DENSE<TLayer, V>.() -> Unit
+    ): Module<T, V> = error(
+        "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " +
+            "dense(outputDimension, dtype = FP16, id) { ... }"
+    )
+
+    @Deprecated(
+        message = "The dtype type argument cannot be implemented (no KClass at runtime, " +
+            "SKEEP-008); pass the dtype as a value: stage(id, dtype = FP16) { ... }",
+        level = DeprecationLevel.ERROR,
+    )
     override fun <TStage : DType> stage(
         id: String,
         content: NeuralNetworkDsl<TStage, V>.() -> Unit
-    ): Module<T, V> {
-        // Create a mixed-precision stage that handles conversion
-        TODO("Mixed-precision stage implementation needed")
-    }
+    ): Module<T, V> = error(
+        "Mixed-precision stage was redesigned (SKEEP-008): pass the dtype as a value, " +
+            "stage(id, dtype = FP16) { ... }"
+    )
 
     override fun batchNorm(
         numFeatures: Int,
@@ -1995,6 +2140,14 @@ public class NeuralNetworkDslImpl<T : DType, V>(
 }
 
 
+/** The [KClass] witness of a dtype instance, typed to the instance's own type (dtypes are objects). */
+@Suppress("UNCHECKED_CAST")
+private fun <T : DType> T.witnessOf(): KClass<T> = this::class as KClass<T>
+
+/** The dtype object whose witness is [kClass]: the DSL impls hold the witness, [PrecisionBoundary] needs the instance. */
+@Suppress("UNCHECKED_CAST")
+private fun <T : DType> dtypeInstanceOf(kClass: KClass<T>): T = DType.fromWitness(kClass) as T
+
 @NetworkDsl
 public class NetworkBuilder<T : DType, V> {
     private val modules = mutableListOf<Module<T, V>>()
diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt
index 7a40dd007..f29a90fc3 100644
--- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt
+++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt
@@ -3,6 +3,10 @@ package sk.ainet.lang.tensor.ops
 import sk.ainet.lang.tensor.Tensor
 import sk.ainet.lang.types.DType
 
+@Deprecated(
+    message = "Never implemented: TensorOps.convert already covers the cast, and the DSL's " +
+        "mixed-precision layers use PrecisionBoundary (SKEEP-008, #1433). Will be removed.",
+)
 public interface MixedPrecisionTensorOps<V> : TensorOps {
     override fun <TFrom : DType, TTo : DType, V> convert(
         tensor: Tensor<TFrom, V>,

From ab7ecbc6e643fa206bcb327b79c080e2c871f8ea Mon Sep 17 00:00:00 2001
From: michalharakal <michalharakal@users.noreply.github.com>
Date: Sun, 1 Nov 2026 11:46:02 +0100
Subject: [PATCH 2/3] fix(cpu): treat BF16 as float-backed across DefaultCpuOps
 and DenseTensorDataFactory

The non-JVM CI lanes failed on MixedPrecisionDslTest because BF16 was missing
from the float-backed dtype branches: DenseTensorDataFactory.init/uniform/
randomInit/wrapFloatArray/adoptFloatArray rejected it, and every DefaultCpuOps
float dispatch (matmul and 33 siblings) plus the sqrt/pow/powScalar/
convTranspose2d/upsample2d guards listed only FP32/FP16. BF16 is float-backed
exactly like FP16, so it joins those branches.
---
 .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 94 ++++++++++---------
 .../tensor/data/DenseTensorDataFactory.kt     | 16 ++--
 2 files changed, 60 insertions(+), 50 deletions(-)

diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt
index 69cd20ef6..f8ed5d828 100644
--- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt
+++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt
@@ -406,7 +406,7 @@ public open class DefaultCpuOpsBase(
         floatBinaryFast(a, b) { x, y -> x + y }?.let { return it }
         return elementwise(a, b) { av, bv, dtype ->
             when (dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val x = av as Float;
                     val y = bv as Float; (x + y) as V
                 }
@@ -429,7 +429,7 @@ public open class DefaultCpuOpsBase(
         floatBinaryFast(a, b) { x, y -> x - y }?.let { return it }
         return elementwise(a, b) { av, bv, dtype ->
             when (dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val x = av as Float;
                     val y = bv as Float; (x - y) as V
                 }
@@ -452,7 +452,7 @@ public open class DefaultCpuOpsBase(
         floatBinaryFast(a, b) { x, y -> x * y }?.let { return it }
         return elementwise(a, b) { av, bv, dtype ->
             when (dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val x = av as Float;
                     val y = bv as Float; (x * y) as V
                 }
@@ -475,7 +475,7 @@ public open class DefaultCpuOpsBase(
         floatBinaryFast(a, b) { x, y -> x / y }?.let { return it }
         return elementwise(a, b) { av, bv, dtype ->
             when (dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val x = av as Float;
                     val y = bv as Float; (x / y) as V
                 }
@@ -886,7 +886,7 @@ public open class DefaultCpuOpsBase(
             val bBatchIdx = mapBatchIndexEff(batchIdx, bEff)
 
             when (a.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     var acc = 0.0f
                     var k = 0
                     while (k < kA) {
@@ -1525,7 +1525,7 @@ public open class DefaultCpuOpsBase(
             val ow = outIdx[3]
 
             when (input.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     var acc = 0.0f
                     val groupIdx = (oc * groups) / cOut
                     val inCStart = groupIdx * cInPerGroup
@@ -1672,7 +1672,7 @@ public open class DefaultCpuOpsBase(
             val ol = outIdx[2]
 
             when (input.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     var acc = 0.0f
                     val groupIdx = (oc * groups) / cOut
                     val inCStart = groupIdx * cInPerGroup
@@ -1830,7 +1830,11 @@ public open class DefaultCpuOpsBase(
         require(groups >= 1) { "convTranspose2d: groups must be >= 1" }
         require(input.dtype == weight.dtype) { "convTranspose2d: dtype mismatch between input and weight: ${input.dtype} vs ${weight.dtype}" }
         bias?.let { require(it.dtype == input.dtype) { "convTranspose2d: dtype mismatch for bias" } }
-        require(input.dtype == sk.ainet.lang.types.FP32::class || input.dtype == sk.ainet.lang.types.FP16::class) {
+        require(
+            input.dtype == sk.ainet.lang.types.FP32::class ||
+                input.dtype == sk.ainet.lang.types.FP16::class ||
+                input.dtype == sk.ainet.lang.types.BF16::class
+        ) {
             "Unsupported dtype for convTranspose2d: ${input.dtype}"
         }
 
@@ -2004,7 +2008,7 @@ public open class DefaultCpuOpsBase(
             val ow = outIdx[4]
 
             when (input.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     var acc = 0.0f
                     val groupIdx = (oc * groups) / cOut
                     val inCStart = groupIdx * cInPerGroup
@@ -2141,7 +2145,7 @@ public open class DefaultCpuOpsBase(
             }
 
             UpsampleMode.Bilinear -> {
-                require(input.dtype == FP32::class || input.dtype == FP16::class) {
+                require(input.dtype == FP32::class || input.dtype == FP16::class || input.dtype == BF16::class) {
                     "upsample2d: Bilinear mode is only implemented for float dtypes (got ${input.dtype})"
                 }
                 dataFactory.init<T, V>(outShape, input.dtype) { idx ->
@@ -2216,7 +2220,7 @@ public open class DefaultCpuOpsBase(
             val hBase = oh * sH - pH
             val wBase = ow * sW - pW
             when (input.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     var best = Float.NEGATIVE_INFINITY
                     var kh = 0
                     while (kh < kH) {
@@ -2295,7 +2299,7 @@ public open class DefaultCpuOpsBase(
             val hBase = oh * sH - pH
             val wBase = ow * sW - pW
             when (input.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     var sum = 0.0f
                     var count = 0
                     var kh = 0
@@ -2625,7 +2629,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> if (x < 0f) 0f else x }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val v = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     (if (v < 0f) 0f else v) as V
@@ -2646,7 +2650,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> if (x < 0f) negativeSlope * x else x }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val v = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     (if (v < 0f) negativeSlope * v else v) as V
@@ -2667,7 +2671,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> if (x >= 0f) x else alpha * (kotlin.math.exp(x) - 1f) }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val v = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     (if (v >= 0f) v else alpha * (kotlin.math.exp(v) - 1f)) as V
@@ -2710,7 +2714,7 @@ public open class DefaultCpuOpsBase(
             }
         }
         when (tensor.dtype) {
-            sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+            sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                 val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { outIdx ->
                     // For the slice defined by outIdx except for nd, compute softmax at that position
                     val idxMax = outIdx.copyOf()
@@ -2771,7 +2775,7 @@ public open class DefaultCpuOpsBase(
             }
         }
         when (tensor.dtype) {
-            sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+            sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                 val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { outIdx ->
                     val idx = outIdx.copyOf()
                     var maxVal = Float.NEGATIVE_INFINITY
@@ -2802,7 +2806,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> 1.0f / (1.0f + kotlin.math.exp(-x)) }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val x = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     (1.0f / (1.0f + kotlin.math.exp(-x))) as V
@@ -2818,7 +2822,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> x * (1.0f / (1.0f + kotlin.math.exp(-x))) }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val x = tensor.data.get(*idx) as Float
                     val s = 1.0f / (1.0f + kotlin.math.exp(-x))
                     @Suppress("UNCHECKED_CAST")
@@ -2840,7 +2844,7 @@ public open class DefaultCpuOpsBase(
         // Tanh approximation of GELU
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val x = tensor.data.get(*idx) as Float
                     val x3 = x * x * x
                     val inner = x + 0.044715f * x3
@@ -2861,7 +2865,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> ErfKernel.geluExact(x) }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val x = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     ErfKernel.geluExact(x) as V
@@ -2877,7 +2881,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> ErfKernel.erf(x.toDouble()).toFloat() }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val x = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     ErfKernel.erf(x.toDouble()).toFloat() as V
@@ -2907,7 +2911,7 @@ public open class DefaultCpuOpsBase(
                 tensor.data.get(*outIdx) as V
             } else {
                 when (tensor.dtype) {
-                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                         @Suppress("UNCHECKED_CAST")
                         0.0f as V
                     }
@@ -2963,7 +2967,7 @@ public open class DefaultCpuOpsBase(
             val outShape = Shape()
             val outData = dataFactory.init<T, V>(outShape, tensor.dtype) {
                 when (tensor.dtype) {
-                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                         var acc = 0.0f
                         val dims = tensor.shape.dimensions
                         if (dims.isEmpty()) {
@@ -3030,7 +3034,7 @@ public open class DefaultCpuOpsBase(
             val outShape = Shape(outDims)
             val outData = dataFactory.init<T, V>(outShape, tensor.dtype) { outIdx ->
                 when (tensor.dtype) {
-                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                         var acc = 0.0f
                         val inIdx = IntArray(rank)
                         var o = 0
@@ -3150,7 +3154,7 @@ public open class DefaultCpuOpsBase(
             val count = if (tensor.volume == 0) 0 else tensor.volume
             val outData = dataFactory.init<T, V>(outShape, tensor.dtype) {
                 when (tensor.dtype) {
-                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                         if (count == 0) {
                             @Suppress("UNCHECKED_CAST") (0.0f as V)
                         } else {
@@ -3218,7 +3222,7 @@ public open class DefaultCpuOpsBase(
             val n = inDims[nd]
             val outData = dataFactory.init<T, V>(outShape, tensor.dtype) { outIdx ->
                 when (tensor.dtype) {
-                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                         if (n == 0) {
                             @Suppress("UNCHECKED_CAST") (0.0f as V)
                         } else {
@@ -3272,7 +3276,7 @@ public open class DefaultCpuOpsBase(
             val n = tensor.volume
             val outData = dataFactory.init<T, V>(outShape, tensor.dtype) {
                 when (tensor.dtype) {
-                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                         if (n == 0) {
                             @Suppress("UNCHECKED_CAST") (0.0f as V)
                         } else {
@@ -3354,7 +3358,7 @@ public open class DefaultCpuOpsBase(
             val n = inDims[nd]
             val outData = dataFactory.init<T, V>(outShape, tensor.dtype) { outIdx ->
                 when (tensor.dtype) {
-                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                         if (n == 0) {
                             @Suppress("UNCHECKED_CAST") (0.0f as V)
                         } else {
@@ -3411,7 +3415,8 @@ public open class DefaultCpuOpsBase(
     override fun <T : DType, V> sqrt(tensor: Tensor<T, V>): Tensor<T, V> {
         require(
             tensor.dtype == sk.ainet.lang.types.FP32::class ||
-                tensor.dtype == sk.ainet.lang.types.FP16::class
+                tensor.dtype == sk.ainet.lang.types.FP16::class ||
+                tensor.dtype == sk.ainet.lang.types.BF16::class
         ) { "sqrt supports only FP16/FP32, got ${tensor.dtype}" }
 
         floatUnaryFast(tensor) { x -> sqrt(x) }?.let { return it }
@@ -3432,7 +3437,8 @@ public open class DefaultCpuOpsBase(
     override fun <T : DType, V> pow(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V> {
         require(
             a.dtype == sk.ainet.lang.types.FP32::class ||
-                a.dtype == sk.ainet.lang.types.FP16::class
+                a.dtype == sk.ainet.lang.types.FP16::class ||
+                a.dtype == sk.ainet.lang.types.BF16::class
         ) { "pow supports only FP16/FP32, got ${a.dtype}" }
         require(a.shape == b.shape) { "pow requires matching shapes; got ${a.shape} and ${b.shape}" }
         floatBinaryFast(a, b) { x, y -> scalarPow(x, y) }?.let { return it }
@@ -3453,7 +3459,8 @@ public open class DefaultCpuOpsBase(
     override fun <T : DType, V> powScalar(a: Tensor<T, V>, n: Number): Tensor<T, V> {
         require(
             a.dtype == sk.ainet.lang.types.FP32::class ||
-                a.dtype == sk.ainet.lang.types.FP16::class
+                a.dtype == sk.ainet.lang.types.FP16::class ||
+                a.dtype == sk.ainet.lang.types.BF16::class
         ) { "powScalar supports only FP16/FP32, got ${a.dtype}" }
         val nFloat = n.toFloat()
         val nInt = n.toInt()
@@ -3493,7 +3500,8 @@ public open class DefaultCpuOpsBase(
     override fun <T : DType, V> log(tensor: Tensor<T, V>): Tensor<T, V> {
         require(
             tensor.dtype == sk.ainet.lang.types.FP32::class ||
-                tensor.dtype == sk.ainet.lang.types.FP16::class
+                tensor.dtype == sk.ainet.lang.types.FP16::class ||
+                tensor.dtype == sk.ainet.lang.types.BF16::class
         ) { "log supports only FP16/FP32, got ${tensor.dtype}" }
         floatUnaryFast(tensor) { x -> ln(x) }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
@@ -3508,7 +3516,8 @@ public open class DefaultCpuOpsBase(
     override fun <T : DType, V> log2(tensor: Tensor<T, V>): Tensor<T, V> {
         require(
             tensor.dtype == sk.ainet.lang.types.FP32::class ||
-                tensor.dtype == sk.ainet.lang.types.FP16::class
+                tensor.dtype == sk.ainet.lang.types.FP16::class ||
+                tensor.dtype == sk.ainet.lang.types.BF16::class
         ) { "log2 supports only FP16/FP32, got ${tensor.dtype}" }
         floatUnaryFast(tensor) { x -> kmLog2(x) }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
@@ -3523,7 +3532,8 @@ public open class DefaultCpuOpsBase(
     override fun <T : DType, V> log10(tensor: Tensor<T, V>): Tensor<T, V> {
         require(
             tensor.dtype == sk.ainet.lang.types.FP32::class ||
-                tensor.dtype == sk.ainet.lang.types.FP16::class
+                tensor.dtype == sk.ainet.lang.types.FP16::class ||
+                tensor.dtype == sk.ainet.lang.types.BF16::class
         ) { "log10 supports only FP16/FP32, got ${tensor.dtype}" }
         floatUnaryFast(tensor) { x -> kmLog10(x) }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
@@ -3541,7 +3551,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> kotlin.math.abs(x) }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val v = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     kotlin.math.abs(v) as V
@@ -3562,7 +3572,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> if (x > 0f) 1f else if (x < 0f) -1f else 0f }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val v = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     (if (v > 0f) 1f else if (v < 0f) -1f else 0f) as V
@@ -3584,7 +3594,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> x.coerceIn(minVal, maxVal) }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val v = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     v.coerceIn(minVal, maxVal) as V
@@ -3605,7 +3615,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> if (x < value) 1f else 0f }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val v = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     (if (v < value) 1f else 0f) as V
@@ -3626,7 +3636,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> if (x >= value) 1f else 0f }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val v = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     (if (v >= value) 1f else 0f) as V
@@ -3677,7 +3687,7 @@ public open class DefaultCpuOpsBase(
             } else {
                 // Zero padding
                 when (tensor.dtype) {
-                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                    sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                         @Suppress("UNCHECKED_CAST") (0f as V)
                     }
                     sk.ainet.lang.types.Int32::class -> {
@@ -3919,7 +3929,7 @@ public open class DefaultCpuOpsBase(
         floatUnaryFast(tensor) { x -> kotlin.math.tanh(x) }?.let { return it }
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
             when (tensor.dtype) {
-                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> {
                     val x = tensor.data.get(*idx) as Float
                     @Suppress("UNCHECKED_CAST")
                     kotlin.math.tanh(x) as V
diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt
index 59e597a46..ff31c71a1 100644
--- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt
+++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt
@@ -521,10 +521,10 @@ public class DenseTensorDataFactory: TensorDataFactory {
                 val data = FloatArray(shape.volume) { random.nextFloat() * range + min }
                 createFloatTensorData(shape, data, FP32 as T) as TensorData<T, V>
             }
-            FP16::class -> {
+            FP16::class, BF16::class -> {
                 val range = max - min
                 val data = FloatArray(shape.volume) { random.nextFloat() * range + min }
-                createFloatTensorData(shape, data, FP16 as T) as TensorData<T, V>
+                createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData<T, V>
             }
             Int32::class -> {
                 val intMin = min.toInt()
@@ -562,7 +562,7 @@ public class DenseTensorDataFactory: TensorDataFactory {
                 }
                 createFloatTensorData(shape, data, FP32 as T) as TensorData<T, V>
             }
-            FP16::class -> {
+            FP16::class, BF16::class -> {
                 val data = FloatArray(shape.volume) { flatIndex ->
                     val indices = IntArray(shape.dimensions.size)
                     var remaining = flatIndex
@@ -572,7 +572,7 @@ public class DenseTensorDataFactory: TensorDataFactory {
                     }
                     (generator(indices) as Float)
                 }
-                createFloatTensorData(shape, data, FP16 as T) as TensorData<T, V>
+                createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData<T, V>
             }
             Int32::class -> {
                 val data = IntArray(shape.volume) { flatIndex ->
@@ -615,9 +615,9 @@ public class DenseTensorDataFactory: TensorDataFactory {
                 val data = FloatArray(shape.volume) { (generator(random) as Float) }
                 createFloatTensorData(shape, data, FP32 as T) as TensorData<T, V>
             }
-            FP16::class -> {
+            FP16::class, BF16::class -> {
                 val data = FloatArray(shape.volume) { (generator(random) as Float) }
-                createFloatTensorData(shape, data, FP16 as T) as TensorData<T, V>
+                createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData<T, V>
             }
             Int32::class -> {
                 val data = IntArray(shape.volume) { (generator(random) as Int) }
@@ -702,7 +702,7 @@ public class DenseTensorDataFactory: TensorDataFactory {
         @Suppress("UNCHECKED_CAST")
         return when (dtype) {
             FP32::class -> DenseFloatArrayTensorData<T>(shape, data) as TensorData<T, V>
-            FP16::class -> DenseFloatArrayTensorData<T>(shape, data) as TensorData<T, V>
+            FP16::class, BF16::class -> DenseFloatArrayTensorData<T>(shape, data) as TensorData<T, V>
             else -> throw IllegalArgumentException("wrapFloatArray only supports floating point types: $dtype")
         }
     }
@@ -714,7 +714,7 @@ public class DenseTensorDataFactory: TensorDataFactory {
         dtype: KClass<T>,
         data: FloatArray
     ): TensorData<T, V> = when (dtype) {
-        FP32::class, FP16::class -> wrapFloatArray(shape, dtype, data)
+        FP32::class, FP16::class, BF16::class -> wrapFloatArray(shape, dtype, data)
         else -> fromFloatArray(shape, dtype, data)
     }
 

From 28bb7c8b5e3643aa0fe50417bac952387d18173e Mon Sep 17 00:00:00 2001
From: michalharakal <michalharakal@users.noreply.github.com>
Date: Sun, 1 Nov 2026 11:58:58 +0100
Subject: [PATCH 3/3] docs: regenerate layer coverage matrix
 (PrecisionBoundary, 33/33)

---
 .../ROOT/pages/reference/operators/generated/index.adoc      | 2 +-
 docs/modules/ROOT/pages/reference/ops-status-matrix.adoc     | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc b/docs/modules/ROOT/pages/reference/operators/generated/index.adoc
index 40f914442..81f7574e4 100644
--- a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc
+++ b/docs/modules/ROOT/pages/reference/operators/generated/index.adoc
@@ -1,6 +1,6 @@
 = AI-NET Operators Reference
 
-Generated from version `0.62.4` on 2026-10-31
+Generated from version `0.62.4` on 2026-11-01
 
 == Operators by Modality
 
diff --git a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc
index e0bf3086c..cbf50c761 100644
--- a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc
+++ b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc
@@ -1,7 +1,7 @@
 = Operator Coverage Matrix
 :description: Cross-backend status for every operator function in SKaiNET.
 
-Generated from `operators.json` version `0.62.4` on 2026-10-31.
+Generated from `operators.json` version `0.62.4` on 2026-11-01.
 
 Rows are `Operator.function` pairs. A gap is marked in source with `@NotImplemented(reason = …, issue = …)` (or `@InProgress`) and shows up here as ❌ (or 🚧) linked to its tracking issue; the PR that closes the gap removes the annotation. See xref:contributing/issue-taxonomy.adoc[Issue taxonomy] to pick one up.
 
@@ -110,6 +110,7 @@ Every concrete `sk.ainet.lang.nn` `Module` subclass. ❌ means the layer's `forw
 | `Linear` | `sk.ainet.lang.nn` | ✅ | —
 | `Lstm` | `sk.ainet.lang.nn` | ✅ | —
 | `MaxPool2d` | `sk.ainet.lang.nn` | ✅ | —
+| `PrecisionBoundary` | `sk.ainet.lang.nn` | ✅ | —
 | `TransposedConv1d` | `sk.ainet.lang.nn` | ✅ | —
 | `TransposedConv2d` | `sk.ainet.lang.nn` | ✅ | —
 | `Upsample2d` | `sk.ainet.lang.nn` | ✅ | —
@@ -129,5 +130,5 @@ Every concrete `sk.ainet.lang.nn` `Module` subclass. ❌ means the layer's `forw
 | `LayerNormalization` | `sk.ainet.lang.nn.normalization` | ✅ | —
 | `MLP` | `sk.ainet.lang.nn.topology` | ✅ | —
 
-| *Done* | | *32 / 32* |
+| *Done* | | *33 / 33* |
 |===