From 6a43302d178d3f1eb7d4e5b81f2d78a72807fd77 Mon Sep 17 00:00:00 2001
From: Benjamin Meyer <benikm91@gmail.com>
Date: Fri, 4 Sep 2026 07:31:43 +0200
Subject: [PATCH 1/5] Prepare 0.1.0 release

Set the version to 0.1.0 and add the Sonatype Central publishing configuration,
mirroring dimwit: credential host, publishTo, licenses, scmInfo, developers and
an early-semver version scheme, plus the sbt-sonatype and sbt-pgp plugins. A
`sonaUploadCore` alias publishes `core` only, since `examples` depends on the
unpublished plotwit.

Depend on the released dimwit-core 0.1.0 rather than 0.2-SNAPSHOT. A release
cannot depend on a snapshot, and this also lets CI drop the temporary step that
cloned dimwit and published it locally. `examples` keeps resolving plotwit from
the local ivy repo, so the dimwit dependency scheme is relaxed there to allow
plotwit's newer snapshot to win.

Document the coordinates in the README install section.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 .github/workflows/ci.yml | 12 ----------
 README.md                | 15 +++++++++++++
 build.sbt                | 47 +++++++++++++++++++++++++++++++++++-----
 mdocs/README.md          | 15 +++++++++++++
 project/plugins.sbt      |  2 ++
 5 files changed, 73 insertions(+), 18 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index ff1d33a..8bf19ab 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -44,18 +44,6 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-sbt-
 
-      # TEMPORARY: dimwit 0.2-SNAPSHOT is not published to any remote yet, so CI builds it from source.
-      - name: Check out dimwit
-        uses: actions/checkout@v4
-        with:
-          repository: marcelluethi/dimwit
-          ref: version-0.2-SNAPSHOT
-          path: dimwit
-
-      - name: Publish dimwit-core to the local ivy repo
-        working-directory: dimwit
-        run: sbt core/publishLocal
-
       - name: Check formatting
         run: sbt scalafmtCheckAll
 
diff --git a/README.md b/README.md
index 896221c..c4a51d3 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,21 @@
 
 DeepWit is a deep learning library for Scala 3 built on one idea: **the code should mirror the theory.** DeepWit is build on [DimWit](https://github.com/dimwit-dev/dimwit), a statically typed tensor library with runtime performance on-par with JAX. DeepWit provides core deep learning modules with their minimal conceptual scope to both express their logic clearly and increase reusablility.
 
+## Installation
+
+DeepWit is published for Scala 3. Add both DeepWit and [DimWit](https://github.com/dimwit-dev/dimwit) —
+DeepWit's API is expressed in DimWit's tensor types, so you will import from both:
+
+```scala
+libraryDependencies ++= Seq(
+  "ch.contrafactus" %% "deepwit-core" % "0.1.0",
+  "ch.contrafactus" %% "dimwit-core" % "0.1.0"
+)
+```
+
+DimWit runs JAX through ScalaPy, so a Python environment with `jax` and `einops` is also required;
+see [DimWit's setup instructions](https://github.com/dimwit-dev/dimwit#installation).
+
 ## Why DeepWit?
 
 DeepWit strips out the framework machinery and leaves the theory standing: no module base class, no indirect loss, no hidden parameters, no hidden gradients, no hidden optimizer states. 
diff --git a/build.sbt b/build.sbt
index e2fe594..10cf179 100644
--- a/build.sbt
+++ b/build.sbt
@@ -4,23 +4,58 @@ import scala.sys.process._
 run / fork := true
 Global / cancelable := true
 
-ThisBuild / version := "0.2-SNAPSHOT"
+ThisBuild / version := "0.1.0"
 ThisBuild / scalaVersion := "3.8.1"
 ThisBuild / organization := "ch.contrafactus"
+ThisBuild / versionScheme := Some("early-semver")
+
+// Publishing to Sonatype Central. The `ch.contrafactus` namespace is verified once for the whole
+// organisation, so deepwit needs no verification of its own.
+ThisBuild / sonatypeCredentialHost := "central.sonatype.com"
+ThisBuild / publishTo := {
+  if (isSnapshot.value)
+    Some("central-snapshots" at "https://central.sonatype.com/repository/maven-snapshots/")
+  else
+    sonatypePublishToBundle.value
+}
+ThisBuild / publishMavenStyle := true
+ThisBuild / homepage := Some(url("https://github.com/dimwit-dev/deepwit"))
+ThisBuild / licenses := List("Apache-2.0" -> url("https://www.apache.org/licenses/LICENSE-2.0"))
+ThisBuild / scmInfo := Some(
+  ScmInfo(
+    url("https://github.com/dimwit-dev/deepwit"),
+    "scm:git@github.com:dimwit-dev/deepwit.git"
+  )
+)
+ThisBuild / developers := List(
+  Developer(
+    id = "dimwit-dev",
+    name = "DeepWit Contributors",
+    email = "",
+    url = url("https://github.com/dimwit-dev")
+  )
+)
 
 // scalafix's RemoveUnused reads the compiler's own unused warnings on Scala 3, so it needs both
 // semanticdb and -Wunused to be on. Run it with `sbt scalafixAll`.
 ThisBuild / semanticdbEnabled := true
 ThisBuild / scalacOptions += "-Wunused:imports"
 
-// deepwit tracks the dimwit/plotwit 0.2 snapshots and is itself snapshot-only for now.
-// NOTE: dimwit-core 0.2-SNAPSHOT is not on this resolver yet (only 0.1.0-SNAPSHOT is), and
-// plotwit is not published at all, so both currently resolve from the local ivy repo via
-// `publishLocal`. CI cannot resolve either until dimwit 0.2-SNAPSHOT is published here.
+// `core` depends only on released artifacts so that it can be published and consumed without any
+// local publishing. `examples` additionally needs plotwit, which is not published anywhere yet and
+// so still resolves from the local ivy repo via `publishLocal` in the plotwit checkout.
 ThisBuild / resolvers += "Central Portal Snapshots" at "https://central.sonatype.com/repository/maven-snapshots/"
 
+// Consequence of that split: `core` asks for the released dimwit 0.1.0 while the locally published
+// plotwit asks for 0.2-SNAPSHOT, which sbt reads as a binary-incompatible conflict under
+// early-semver. Let the newer snapshot win in `examples`; `core` on its own resolves 0.1.0.
+ThisBuild / libraryDependencySchemes += "ch.contrafactus" %% "dimwit-core" % "always"
+
 addCommandAlias("testAndCoverage", "; clean; coverage; test; coverageReport")
 
+// Publishes `core` only. `examples` depends on plotwit, which is not published.
+addCommandAlias("sonaUploadCore", "; project core; sonatypeCentralUpload; project root")
+
 lazy val uvPython: String =
   sys.env.getOrElse(
     "DIMWIT_PYTHON_PATH",
@@ -47,7 +82,7 @@ lazy val core = (project in file("core"))
       "org.scalacheck" %% "scalacheck" % "1.18.0" % Test,
       "org.scalatestplus" %% "scalacheck-1-18" % "3.2.19.0" % Test,
       "dev.scalapy" %% "scalapy-core" % "0.5.3",
-      "ch.contrafactus" %% "dimwit-core" % "0.2-SNAPSHOT" changing ()
+      "ch.contrafactus" %% "dimwit-core" % "0.1.0"
     ),
     // ScalaPy drives a single embedded CPython interpreter, and two suites importing jax at the same
     // time race into a partially initialized module. Whichever suites happen to touch a tensor first
diff --git a/mdocs/README.md b/mdocs/README.md
index 7e61d83..3825016 100644
--- a/mdocs/README.md
+++ b/mdocs/README.md
@@ -22,6 +22,21 @@ trait Embedding derives Label
 
 DeepWit is a deep learning library for Scala 3 built on one idea: **the code should mirror the theory.** DeepWit is build on [DimWit](https://github.com/dimwit-dev/dimwit), a statically typed tensor library with runtime performance on-par with JAX. DeepWit provides core deep learning modules with their minimal conceptual scope to both express their logic clearly and increase reusablility.
 
+## Installation
+
+DeepWit is published for Scala 3. Add both DeepWit and [DimWit](https://github.com/dimwit-dev/dimwit) —
+DeepWit's API is expressed in DimWit's tensor types, so you will import from both:
+
+```scala
+libraryDependencies ++= Seq(
+  "ch.contrafactus" %% "deepwit-core" % "@VERSION@",
+  "ch.contrafactus" %% "dimwit-core" % "0.1.0"
+)
+```
+
+DimWit runs JAX through ScalaPy, so a Python environment with `jax` and `einops` is also required;
+see [DimWit's setup instructions](https://github.com/dimwit-dev/dimwit#installation).
+
 ## Why DeepWit?
 
 DeepWit strips out the framework machinery and leaves the theory standing: no module base class, no indirect loss, no hidden parameters, no hidden gradients, no hidden optimizer states. 
diff --git a/project/plugins.sbt b/project/plugins.sbt
index d0a3ad4..3dd7a80 100644
--- a/project/plugins.sbt
+++ b/project/plugins.sbt
@@ -2,5 +2,7 @@ addSbtPlugin("org.scalameta" % "sbt-scalafmt" % "2.5.5")
 addSbtPlugin("org.scalameta" % "sbt-mdoc" % "2.8.2")
 addSbtPlugin("ch.epfl.scala" % "sbt-scalafix" % "0.14.7")
 addSbtPlugin("org.scoverage" % "sbt-scoverage" % "2.3.0")
+addSbtPlugin("org.xerial.sbt" % "sbt-sonatype" % "3.11.3")
+addSbtPlugin("com.github.sbt" % "sbt-pgp" % "2.3.1")
 
 libraryDependencies += "ai.kien" %% "python-native-libs" % "0.2.2"

From 87af391e57a2815b5a4327c74c959bf7618f0095 Mon Sep 17 00:00:00 2001
From: Benjamin Meyer <benikm91@gmail.com>
Date: Fri, 4 Sep 2026 07:57:58 +0200
Subject: [PATCH 2/5] Set a real POM description and stop publishing test
 artifacts

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 build.sbt | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/build.sbt b/build.sbt
index 10cf179..8a6ffdb 100644
--- a/build.sbt
+++ b/build.sbt
@@ -92,8 +92,12 @@ lazy val core = (project in file("core"))
     // SCALAPY_PYTHON_LIBRARY / SCALAPY_PYTHON_PROGRAMNAME being exported by the shell.
     fork := true,
     javaOptions ++= scalapyJavaOptions,
+    description := "A theory-aligned deep learning library for Scala 3, built on DimWit",
     Compile / packageSrc / publishArtifact := true,
-    Compile / packageDoc / publishArtifact := true
+    Compile / packageDoc / publishArtifact := true,
+    // Ship the library's own sources and docs, not the test ones.
+    Test / packageSrc / publishArtifact := false,
+    Test / packageDoc / publishArtifact := false
   )
 
 // Examples subproject

From d66f05926a3534c66ad93825db7578a986428632 Mon Sep 17 00:00:00 2001
From: Benjamin Meyer <benikm91@gmail.com>
Date: Fri, 4 Sep 2026 08:09:30 +0200
Subject: [PATCH 3/5] Drop the raw-JAX escape in VocabularyEmbedder and make
 Huber's transition point traceable
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

VocabularyEmbedder reached past DimWit into `Jax.jnp.take` to look up an
embedding row. `slice(Axis[Vocab].at(token))` expresses the same thing: the
`AxisAtTensorIndex` overload takes a `Tensor0[Int32]` index and removes the
sliced axis, giving `Tensor1[Embedding, V]` directly. The `take` alternative
left in a comment could not have worked — it needs a `Tensor1` of indices and
replaces an axis rather than removing it.

Rename Huber's `threshold` to `transitionPoint`, which is what the error message
and the tests already called it, and take it as a `Tensor0[Float32]` so it can be
traced and scheduled, matching `Perturbation.thin`. The eager `require` goes with
it: a traced scalar cannot be validated host-side, so the constraint is stated in
the scaladoc instead, as `Perturbation` already does for its probability.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 .../deepwit/embedder/VocabularyEmbedder.scala     |  7 +------
 core/src/main/scala/deepwit/loss/Regression.scala | 12 ++++++++----
 .../test/scala/deepwit/loss/RegressionSuite.scala | 15 ++++++---------
 3 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala b/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
index e55ba16..b30e447 100644
--- a/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
+++ b/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
@@ -3,17 +3,12 @@ package deepwit.embedder
 import dimwit.*
 import dimwit.stats.Normal
 import dimwit.stats.Uniform
-import dimwit.jax.Jax
-import dimwit.python.PyBridge.{toPyTensor, liftPyTensor}
 import dimwit.Label as Λ
 
 class VocabularyEmbedder[Vocab: Λ, Embedding: Λ, V: IsFloating](params: VocabularyEmbedder.Params[Vocab, Embedding, V]) extends (Tensor0[Int32] => Tensor1[Embedding, V]):
 
   override def apply(token: Tensor0[Int32]): Tensor1[Embedding, V] =
-    // params.vocabularyEmbeddings.slice(Axis[Vocab].at(token)) // TODO
-    // params.vocabularyEmbeddings.take(Axis[Vocab])(token)
-    val rawJax = Jax.jnp.take(toPyTensor(params.vocabularyEmbeddings), toPyTensor(token), axis = 0)
-    liftPyTensor(rawJax)
+    params.vocabularyEmbeddings.slice(Axis[Vocab].at(token))
 
   def unembed(embedding: Tensor1[Embedding, V]): Tensor1[Vocab, V] =
     embedding.dot(Axis[Embedding])(params.vocabularyEmbeddings)
diff --git a/core/src/main/scala/deepwit/loss/Regression.scala b/core/src/main/scala/deepwit/loss/Regression.scala
index 0b2ece4..0b96214 100644
--- a/core/src/main/scala/deepwit/loss/Regression.scala
+++ b/core/src/main/scala/deepwit/loss/Regression.scala
@@ -15,9 +15,13 @@ object AbsoluteError:
 
 object Huber:
 
-  def apply[V: IsFloating](target: Tensor0[V], prediction: Tensor0[V], threshold: Float): Tensor0[V] =
-    require(threshold > 0f, s"A transition point must be positive, but was $threshold.")
+  /** @param transitionPoint Where the quadratic branch hands over to the linear one. Must be
+    *                        positive. A scalar rather than a `Float` so it can be traced, and
+    *                        scheduled the way a learning rate is.
+    */
+  def apply[V: IsFloating](target: Tensor0[V], prediction: Tensor0[V], transitionPoint: Tensor0[Float32]): Tensor0[V] =
+    val δ = transitionPoint.asFloat(VType[V])
     val residual = AbsoluteError(target, prediction)
     val squared = 0.5f * SquaredError(target, prediction)
-    val absolute = threshold * (residual - 0.5f * threshold)
-    where(residual <= threshold, squared, absolute)
+    val absolute = δ * (residual - 0.5f * δ)
+    where(residual <= δ, squared, absolute)
diff --git a/core/src/test/scala/deepwit/loss/RegressionSuite.scala b/core/src/test/scala/deepwit/loss/RegressionSuite.scala
index 04de571..5b89102 100644
--- a/core/src/test/scala/deepwit/loss/RegressionSuite.scala
+++ b/core/src/test/scala/deepwit/loss/RegressionSuite.scala
@@ -41,19 +41,16 @@ class RegressionSuite extends AnyFunSpec with Matchers:
   describe("Huber"):
 
     it("matches the halved squared error within the transition point"):
-      val loss = Huber(Tensor0(0f), Tensor0(0.5f), threshold = 1f)
+      val loss = Huber(Tensor0(0f), Tensor0(0.5f), transitionPoint = Tensor0(1f))
       loss.item shouldBe (0.5f * 0.25f +- 1e-6f)
 
     it("grows linearly beyond the transition point"):
-      val atFive = Huber(Tensor0(0f), Tensor0(5f), threshold = 1f).item
-      val atSix = Huber(Tensor0(0f), Tensor0(6f), threshold = 1f).item
+      val atFive = Huber(Tensor0(0f), Tensor0(5f), transitionPoint = Tensor0(1f)).item
+      val atSix = Huber(Tensor0(0f), Tensor0(6f), transitionPoint = Tensor0(1f)).item
       atSix - atFive shouldBe (1f +- 1e-4f)
 
     it("joins its two branches continuously at the transition point"):
-      val threshold = 2f
-      val below = Huber(Tensor0(0f), Tensor0(threshold - 1e-3f), threshold).item
-      val above = Huber(Tensor0(0f), Tensor0(threshold + 1e-3f), threshold).item
+      val transitionPoint = 2f
+      val below = Huber(Tensor0(0f), Tensor0(transitionPoint - 1e-3f), Tensor0(transitionPoint)).item
+      val above = Huber(Tensor0(0f), Tensor0(transitionPoint + 1e-3f), Tensor0(transitionPoint)).item
       below shouldBe (above +- 1e-2f)
-
-    it("rejects a non-positive transition point"):
-      an[IllegalArgumentException] should be thrownBy Huber(Tensor0(0f), Tensor0(1f), threshold = 0f)

From cf1ab1b7d0c2df1ba197647bc374e48ff8c12ef4 Mon Sep 17 00:00:00 2001
From: Benjamin Meyer <benikm91@gmail.com>
Date: Fri, 4 Sep 2026 08:15:58 +0200
Subject: [PATCH 4/5] Keep Huber's transition point a Float so it stays
 validated
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Reverts the Tensor0 change: a traced scalar cannot be checked host-side, and
losing the `require` costs more here than gaining a schedulable transition point.
The rename to `transitionPoint` stays — it is what the error message and the
tests already called it.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 core/src/main/scala/deepwit/loss/Regression.scala   | 12 ++++--------
 .../test/scala/deepwit/loss/RegressionSuite.scala   | 13 ++++++++-----
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/core/src/main/scala/deepwit/loss/Regression.scala b/core/src/main/scala/deepwit/loss/Regression.scala
index 0b96214..cd16f4c 100644
--- a/core/src/main/scala/deepwit/loss/Regression.scala
+++ b/core/src/main/scala/deepwit/loss/Regression.scala
@@ -15,13 +15,9 @@ object AbsoluteError:
 
 object Huber:
 
-  /** @param transitionPoint Where the quadratic branch hands over to the linear one. Must be
-    *                        positive. A scalar rather than a `Float` so it can be traced, and
-    *                        scheduled the way a learning rate is.
-    */
-  def apply[V: IsFloating](target: Tensor0[V], prediction: Tensor0[V], transitionPoint: Tensor0[Float32]): Tensor0[V] =
-    val δ = transitionPoint.asFloat(VType[V])
+  def apply[V: IsFloating](target: Tensor0[V], prediction: Tensor0[V], transitionPoint: Float): Tensor0[V] =
+    require(transitionPoint > 0f, s"A transition point must be positive, but was $transitionPoint.")
     val residual = AbsoluteError(target, prediction)
     val squared = 0.5f * SquaredError(target, prediction)
-    val absolute = δ * (residual - 0.5f * δ)
-    where(residual <= δ, squared, absolute)
+    val absolute = transitionPoint * (residual - 0.5f * transitionPoint)
+    where(residual <= transitionPoint, squared, absolute)
diff --git a/core/src/test/scala/deepwit/loss/RegressionSuite.scala b/core/src/test/scala/deepwit/loss/RegressionSuite.scala
index 5b89102..683fcc9 100644
--- a/core/src/test/scala/deepwit/loss/RegressionSuite.scala
+++ b/core/src/test/scala/deepwit/loss/RegressionSuite.scala
@@ -41,16 +41,19 @@ class RegressionSuite extends AnyFunSpec with Matchers:
   describe("Huber"):
 
     it("matches the halved squared error within the transition point"):
-      val loss = Huber(Tensor0(0f), Tensor0(0.5f), transitionPoint = Tensor0(1f))
+      val loss = Huber(Tensor0(0f), Tensor0(0.5f), transitionPoint = 1f)
       loss.item shouldBe (0.5f * 0.25f +- 1e-6f)
 
     it("grows linearly beyond the transition point"):
-      val atFive = Huber(Tensor0(0f), Tensor0(5f), transitionPoint = Tensor0(1f)).item
-      val atSix = Huber(Tensor0(0f), Tensor0(6f), transitionPoint = Tensor0(1f)).item
+      val atFive = Huber(Tensor0(0f), Tensor0(5f), transitionPoint = 1f).item
+      val atSix = Huber(Tensor0(0f), Tensor0(6f), transitionPoint = 1f).item
       atSix - atFive shouldBe (1f +- 1e-4f)
 
     it("joins its two branches continuously at the transition point"):
       val transitionPoint = 2f
-      val below = Huber(Tensor0(0f), Tensor0(transitionPoint - 1e-3f), Tensor0(transitionPoint)).item
-      val above = Huber(Tensor0(0f), Tensor0(transitionPoint + 1e-3f), Tensor0(transitionPoint)).item
+      val below = Huber(Tensor0(0f), Tensor0(transitionPoint - 1e-3f), transitionPoint).item
+      val above = Huber(Tensor0(0f), Tensor0(transitionPoint + 1e-3f), transitionPoint).item
       below shouldBe (above +- 1e-2f)
+
+    it("rejects a non-positive transition point"):
+      an[IllegalArgumentException] should be thrownBy Huber(Tensor0(0f), Tensor0(1f), transitionPoint = 0f)

From 6e19b147373d27a9f629b91ca896dfad9d0d50c9 Mon Sep 17 00:00:00 2001
From: Benjamin Meyer <benikm91@gmail.com>
Date: Fri, 4 Sep 2026 09:54:49 +0200
Subject: [PATCH 5/5] Add comments. Add ReferenceMultiHeadAttentionSuite.

---
 .../ReferenceMultiHeadAttention.scala         |  9 ++-
 .../scala/deepwit/attention/package.scala     |  8 ++-
 .../scala/deepwit/cnn/AffineConv2DLayer.scala |  5 ++
 .../scala/deepwit/cnn/LinearConv2DLayer.scala |  5 ++
 .../scala/deepwit/cnn/MaxPool2DLayer.scala    |  8 ++-
 .../cnn/TransposeAffineConv2DLayer.scala      |  8 +++
 .../cnn/TransposeLinearConv2DLayer.scala      |  8 +++
 .../deepwit/embedder/PositionalEncoding.scala | 12 +++-
 .../deepwit/embedder/VocabularyEmbedder.scala |  7 +++
 .../scala/deepwit/init/Initialization.scala   | 11 +++-
 .../main/scala/deepwit/loss/Regression.scala  |  4 ++
 .../deepwit/normalization/LayerNorm.scala     |  5 ++
 .../scala/deepwit/normalization/RMSNorm.scala |  6 +-
 .../optimizer/LearningRateSchedule.scala      | 37 ++++++------
 .../ReferenceMultiHeadAttentionSuite.scala    | 56 +++++++++++++++++++
 15 files changed, 159 insertions(+), 30 deletions(-)
 create mode 100644 core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala

diff --git a/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala b/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala
index 8f4ed16..ae72e19 100644
--- a/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala
+++ b/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala
@@ -4,10 +4,13 @@ import dimwit.*
 import deepwit.base.AffineLayer
 import dimwit.Label as Λ
 
-/** TODO
-  * This implementation is conceptually clearer than [[MultiHeadAttention]] but slower. How to merge?
+/** A readable statement of what [[MultiHeadAttention]] computes, kept for reference rather than for
+  * use: each head is its own [[Attention]] in a `List`, so the per-head structure stays visible
+  * instead of being folded into batched tensors. That is also what makes it slower — the per-head
+  * parameters are separate tensors rather than one.
+  *
+  * [[MultiHeadAttention]] is tested against it, so the batched implementation is known to agree.
   */
-
 class ReferenceMultiHeadAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding: Λ, V: IsFloating](
     params: ReferenceMultiHeadAttention.Params[SourceEmbedding, TargetEmbedding, V],
     createAttentionMask: Shape2[Target, Source] => Tensor2[Target, Source, Bool],
diff --git a/core/src/main/scala/deepwit/attention/package.scala b/core/src/main/scala/deepwit/attention/package.scala
index 1efbb5d..20e503c 100644
--- a/core/src/main/scala/deepwit/attention/package.scala
+++ b/core/src/main/scala/deepwit/attention/package.scala
@@ -2,8 +2,14 @@ package deepwit.attention
 
 import dimwit.Label
 
-/** Axis labels for the per-head spaces of a [[MultiHeadAttention]]. */
+/** The multi-head attention heads. */
 trait Head derives Label
+
+/** The space a head projects queries into. */
 trait HeadQuery derives Label
+
+/** The space a head projects keys into. */
 trait HeadKey derives Label
+
+/** The space a head projects values into. */
 trait HeadValue derives Label
diff --git a/core/src/main/scala/deepwit/cnn/AffineConv2DLayer.scala b/core/src/main/scala/deepwit/cnn/AffineConv2DLayer.scala
index 0f8ed9a..707b602 100644
--- a/core/src/main/scala/deepwit/cnn/AffineConv2DLayer.scala
+++ b/core/src/main/scala/deepwit/cnn/AffineConv2DLayer.scala
@@ -3,6 +3,11 @@ package deepwit.cnn
 import dimwit.*
 import dimwit.Label as Λ
 
+/** A 2D convolution with bias. For the bias-free equivalent, see [[LinearConv2DLayer]].
+  *
+  * @param stride An `Int` stride applies to both spatial axes.
+  * @param padding `Padding.SAME` preserves the spatial extents, while `Padding.VALID` shrinks them.
+  */
 class AffineConv2DLayer[S1: Λ, S2: Λ, InChannel: Λ, OutChannel: Λ, V: IsFloating](
     params: AffineConv2DLayer.Params[S1, S2, InChannel, OutChannel, V],
     stride: Stride2[S1, S2] | Int = 1,
diff --git a/core/src/main/scala/deepwit/cnn/LinearConv2DLayer.scala b/core/src/main/scala/deepwit/cnn/LinearConv2DLayer.scala
index 0454b92..2665373 100644
--- a/core/src/main/scala/deepwit/cnn/LinearConv2DLayer.scala
+++ b/core/src/main/scala/deepwit/cnn/LinearConv2DLayer.scala
@@ -3,6 +3,11 @@ package deepwit.cnn
 import dimwit.*
 import dimwit.Label as Λ
 
+/** A 2D convolution without bias. For the bias-carrying equivalent, see [[AffineConv2DLayer]].
+  *
+  * @params stride An `Int` stride applies to both spatial axes.
+  * @params padding `Padding.SAME` preserves the spatial extents, while `Padding.VALID` shrinks them.
+  */
 class LinearConv2DLayer[S1: Λ, S2: Λ, InChannel: Λ, OutChannel: Λ, V: IsFloating](
     params: LinearConv2DLayer.Params[S1, S2, InChannel, OutChannel, V],
     stride: Stride2[S1, S2] | Int = 1,
diff --git a/core/src/main/scala/deepwit/cnn/MaxPool2DLayer.scala b/core/src/main/scala/deepwit/cnn/MaxPool2DLayer.scala
index cebad27..2ff8ad6 100644
--- a/core/src/main/scala/deepwit/cnn/MaxPool2DLayer.scala
+++ b/core/src/main/scala/deepwit/cnn/MaxPool2DLayer.scala
@@ -5,9 +5,15 @@ import dimwit.jax.Jax
 import dimwit.python.PyBridge.{liftPyTensor, toPyTensor}
 import dimwit.Label as Λ
 
+/** A sliding-window maximum over the two spatial axes.
+  *
+  * @param window An `Int` window applies to both spatial axes.
+  * @param stride An `Int` stride applies to both spatial axes.
+  * @param padding `Padding.SAME` preserves the spatial extents, while `Padding.VALID` shrinks them.
+  */
 class MaxPool2DLayer[S1: Λ, S2: Λ, V: IsFloating](
     window: Window2[S1, S2] | Int,
-    stride: Stride2[S1, S2] | Int = 1,
+    stride: Stride2[S1, S2] | Int,
     padding: Padding = Padding.SAME
 ) extends (Tensor2[S1, S2, V] => Tensor2[S1, S2, V]):
 
diff --git a/core/src/main/scala/deepwit/cnn/TransposeAffineConv2DLayer.scala b/core/src/main/scala/deepwit/cnn/TransposeAffineConv2DLayer.scala
index 23a8791..7a9b0c0 100644
--- a/core/src/main/scala/deepwit/cnn/TransposeAffineConv2DLayer.scala
+++ b/core/src/main/scala/deepwit/cnn/TransposeAffineConv2DLayer.scala
@@ -3,6 +3,14 @@ package deepwit.cnn
 import dimwit.*
 import dimwit.Label as Λ
 
+/** The adjoint of a 2D convolution, with bias. For the bias-free equivalent, see
+  * [[TransposeLinearConv2DLayer]].
+  *
+  * @tparam InChannel The forward convolution's input channels, so this layer's output.
+  * @tparam OutChannel The forward convolution's output channels, so this layer's input.
+  * @param stride An `Int` stride applies to both spatial axes. A stride above 1 grows the spatial extents, which is how these layers upsample.
+  * @param padding `Padding.SAME` preserves the spatial extents, while `Padding.VALID` grows them.
+  */
 class TransposeAffineConv2DLayer[S1: Λ, S2: Λ, InChannel: Λ, OutChannel: Λ, V: IsFloating](
     params: TransposeAffineConv2DLayer.Params[S1, S2, InChannel, OutChannel, V],
     stride: Stride2[S1, S2] | Int = 1,
diff --git a/core/src/main/scala/deepwit/cnn/TransposeLinearConv2DLayer.scala b/core/src/main/scala/deepwit/cnn/TransposeLinearConv2DLayer.scala
index 0742081..a738169 100644
--- a/core/src/main/scala/deepwit/cnn/TransposeLinearConv2DLayer.scala
+++ b/core/src/main/scala/deepwit/cnn/TransposeLinearConv2DLayer.scala
@@ -3,6 +3,14 @@ package deepwit.cnn
 import dimwit.*
 import dimwit.Label as Λ
 
+/** The adjoint of a 2D convolution, without bias. For the bias-carrying equivalent, see
+  * [[TransposeAffineConv2DLayer]].
+  *
+  * @tparam InChannel The forward convolution's input channels, so this layer's output.
+  * @tparam OutChannel The forward convolution's output channels, so this layer's input.
+  * @param stride An `Int` stride applies to both spatial axes. A stride above 1 grows the spatial extents, which is how these layers upsample.
+  * @param padding `Padding.SAME` preserves the spatial extents, while `Padding.VALID` grows them.
+  */
 class TransposeLinearConv2DLayer[S1: Λ, S2: Λ, InChannel: Λ, OutChannel: Λ, V: IsFloating](
     params: TransposeLinearConv2DLayer.Params[S1, S2, InChannel, OutChannel, V],
     stride: Stride2[S1, S2] | Int = 1,
diff --git a/core/src/main/scala/deepwit/embedder/PositionalEncoding.scala b/core/src/main/scala/deepwit/embedder/PositionalEncoding.scala
index f2a8ab1..62be2b8 100644
--- a/core/src/main/scala/deepwit/embedder/PositionalEncoding.scala
+++ b/core/src/main/scala/deepwit/embedder/PositionalEncoding.scala
@@ -6,7 +6,7 @@ import dimwit.Label as Λ
 
 object PositionalEncoding:
 
-  /** The axis the frequencies live on, half of which become sines and half cosines. */
+  /** The axis the frequencies live on, each contributing both a sine and a cosine. */
   private trait Scale derives Label
 
   /** The default ratio between the fastest and the slowest oscillation, as chosen for sequences of
@@ -18,11 +18,13 @@ object PositionalEncoding:
   def gridPositions[P: Λ, V: IsFloating](extent: AxisExtent[P], vtype: VType[V]): Tensor1[P, V] =
     Tensor1(Axis[P]).fromArray(Array.range(0, extent.size)).asFloat(vtype)
 
-  /** Encodes each position as sines and cosines of it, at geometrically spaced frequencies.
+  /** Encodes each position as sines and cosines of it, at geometrically spaced frequencies, as
+    * described in [Attention Is All You Need](https://arxiv.org/abs/1706.03762).
     *
     * The positions are given rather than derived, so they need not be a grid's indices: any
     * position the encoding is evaluated at means the same thing to a model trained on any other.
     *
+    * @param embeddingExtent Must be even, to pair each frequency's sine with its cosine.
     * @param frequencyRange The ratio between the fastest and the slowest oscillation.
     */
   def sinusoidal[P: Λ, Embedding: Λ, V: IsFloating](
@@ -41,7 +43,11 @@ object PositionalEncoding:
     val scaled = positions.vmap(Axis[P])(_ *! scales)
     concatenate(scaled.sin, scaled.cos, concatAxis = Axis[Scale]).relabel(Axis[Scale], Axis[Embedding])
 
-  /** Encodes each point of a plane, giving each of the two axes half of the embedding. */
+  /** Encodes each point of a plane, giving each of the two axes half of the embedding.
+    *
+    * @param embeddingExtent Must be divisible by four: each axis takes half, and each half pairs sines with cosines.
+    * @param frequencyRange The ratio between the fastest and the slowest oscillation.
+    */
   def sinusoidal2D[X: Λ, Y: Λ, Embedding: Λ, V: IsFloating](
       xPositions: Tensor1[X, V],
       yPositions: Tensor1[Y, V],
diff --git a/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala b/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
index b30e447..a425a1a 100644
--- a/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
+++ b/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
@@ -5,11 +5,15 @@ import dimwit.stats.Normal
 import dimwit.stats.Uniform
 import dimwit.Label as Λ
 
+/** Maps a token to its embedding. [[VocabularyEmbedder.unembed]] maps back (weight typing). */
 class VocabularyEmbedder[Vocab: Λ, Embedding: Λ, V: IsFloating](params: VocabularyEmbedder.Params[Vocab, Embedding, V]) extends (Tensor0[Int32] => Tensor1[Embedding, V]):
 
   override def apply(token: Tensor0[Int32]): Tensor1[Embedding, V] =
     params.vocabularyEmbeddings.slice(Axis[Vocab].at(token))
 
+  /** Scores every token by projecting onto its embedding, through the matrix the lookup uses —
+    * weight tying, as described in [Using the Output Embedding to Improve Language Models](https://arxiv.org/abs/1608.05859).
+    */
   def unembed(embedding: Tensor1[Embedding, V]): Tensor1[Vocab, V] =
     embedding.dot(Axis[Embedding])(params.vocabularyEmbeddings)
 
@@ -17,6 +21,9 @@ object VocabularyEmbedder:
 
   case class Params[Vocab, Embedding, V](vocabularyEmbeddings: Tensor2[Vocab, Embedding, V])
 
+  /** Scaled by the embedding size — a lookup has no fan-in of its own, and this leaves each row
+    * near unit norm.
+    */
   object Params:
 
     def init[Vocab: Λ, Embedding: Λ, V: IsFloating](vocabExtent: AxisExtent[Vocab], embeddingExtent: AxisExtent[Embedding], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1.0): Params[Vocab, Embedding, V] =
diff --git a/core/src/main/scala/deepwit/init/Initialization.scala b/core/src/main/scala/deepwit/init/Initialization.scala
index 73f1c08..a4054d8 100644
--- a/core/src/main/scala/deepwit/init/Initialization.scala
+++ b/core/src/main/scala/deepwit/init/Initialization.scala
@@ -5,6 +5,9 @@ import dimwit.stats.Normal
 import dimwit.stats.Uniform
 import dimwit.Label as Λ
 
+/** Xavier/Glorot initializers, drawing weights with variance `2 / (fanIn + fanOut)` as described in
+  * [Understanding the difficulty of training deep feedforward neural networks](https://proceedings.mlr.press/v9/glorot10a.html).
+  */
 object Init:
 
   def xavierNormal[FanIn: Λ, FanOut: Λ, V: IsFloating](fanIn: AxisExtent[FanIn], fanOut: AxisExtent[FanOut], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1f): Tensor2[FanIn, FanOut, V] =
@@ -13,14 +16,18 @@ object Init:
 
   def xavierUniform[FanIn: Λ, FanOut: Λ, V: IsFloating](fanIn: AxisExtent[FanIn], fanOut: AxisExtent[FanOut], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1f): Tensor2[FanIn, FanOut, V] =
     val variance = Tensor0(vtype)(2.0f / (fanIn.size + fanOut.size))
+    // A uniform on [-a, a] has variance a²/3, so a = √(3·variance) hits the target spread.
     val a = gain * (3f * variance).sqrt
     IndependentDistribution.fromUnivariate(Shape(fanIn, fanOut), Uniform(-a, a)).sample(key)
 
   def xavierNormalVector[FanIn: Λ, V: IsFloating](fanIn: AxisExtent[FanIn], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1f): Tensor1[FanIn, V] =
-    val variance = Tensor0(vtype)(2.0f / (fanIn.size + 1))
+    val fanOut = 1 // linear form has fan-out of 1
+    val variance = Tensor0(vtype)(2.0f / (fanIn.size + fanOut))
     Normal.standardIsotropic(Shape(fanIn), scale = gain * variance.sqrt).sample(key)
 
   def xavierUniformVector[FanIn: Λ, V: IsFloating](fanIn: AxisExtent[FanIn], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1f): Tensor1[FanIn, V] =
-    val variance = Tensor0(vtype)(2.0f / (fanIn.size + 1))
+    val fanOut = 1 // linear form has fan-out of 1
+    val variance = Tensor0(vtype)(2.0f / (fanIn.size + fanOut))
+    // A uniform on [-a, a] has variance a²/3, so a = √(3·variance) hits the target spread.
     val a = gain * (3f * variance).sqrt
     IndependentDistribution.fromUnivariate(Shape(fanIn), Uniform(-a, a)).sample(key)
diff --git a/core/src/main/scala/deepwit/loss/Regression.scala b/core/src/main/scala/deepwit/loss/Regression.scala
index cd16f4c..3d1ad59 100644
--- a/core/src/main/scala/deepwit/loss/Regression.scala
+++ b/core/src/main/scala/deepwit/loss/Regression.scala
@@ -15,9 +15,13 @@ object AbsoluteError:
 
 object Huber:
 
+  /** Quadratic like [[SquaredError]] if residual within `transitionPoint`, linear like [[AbsoluteError]] beyond it as described in
+    * [Robust Estimation of a Location Parameter](https://doi.org/10.1214/aoms/1177703732).
+    */
   def apply[V: IsFloating](target: Tensor0[V], prediction: Tensor0[V], transitionPoint: Float): Tensor0[V] =
     require(transitionPoint > 0f, s"A transition point must be positive, but was $transitionPoint.")
     val residual = AbsoluteError(target, prediction)
+    // Scale squared and absolute errors to meet in value and slope at the transition point.
     val squared = 0.5f * SquaredError(target, prediction)
     val absolute = transitionPoint * (residual - 0.5f * transitionPoint)
     where(residual <= transitionPoint, squared, absolute)
diff --git a/core/src/main/scala/deepwit/normalization/LayerNorm.scala b/core/src/main/scala/deepwit/normalization/LayerNorm.scala
index 0efe825..aec5ba8 100644
--- a/core/src/main/scala/deepwit/normalization/LayerNorm.scala
+++ b/core/src/main/scala/deepwit/normalization/LayerNorm.scala
@@ -6,6 +6,11 @@ import dimwit.Label as Λ
 
 import deepwit.{defaultEpsilon, unwrapEpsilon}
 
+/** Standardizes over the `L` axis, then scales and shifts by the learned parameters, as described
+  * in [Layer Normalization](https://arxiv.org/abs/1607.06450).
+  *
+  * @param epsilon Guards the division. Defaults to the machine epsilon of data type; pass a `Float` to fix it. Pass a function to derive it from the data type.
+  */
 class LayerNorm[L: Λ, V: IsFloating](
     params: LayerNorm.Params[L, V],
     epsilon: Float | (DType => Float) = defaultEpsilon
diff --git a/core/src/main/scala/deepwit/normalization/RMSNorm.scala b/core/src/main/scala/deepwit/normalization/RMSNorm.scala
index 4f2964c..78c81ec 100644
--- a/core/src/main/scala/deepwit/normalization/RMSNorm.scala
+++ b/core/src/main/scala/deepwit/normalization/RMSNorm.scala
@@ -6,6 +6,11 @@ import dimwit.Label as Λ
 
 import deepwit.{defaultEpsilon, unwrapEpsilon}
 
+/** Rescales by the root mean square over the `L` axis, then scales by the learned weight, as
+  * described in [Root Mean Square Layer Normalization](https://arxiv.org/abs/1910.07467).
+  *
+  * @param epsilon Guards the division. Defaults to the machine epsilon of data type; pass a `Float` to fix it. Pass a function to derive it from the data type.
+  */
 class RMSNorm[L: Λ, V: IsFloating](
     params: RMSNorm.Params[L, V],
     epsilon: Float | (DType => Float) = defaultEpsilon
@@ -15,7 +20,6 @@ class RMSNorm[L: Λ, V: IsFloating](
 
   def apply(x: Tensor1[L, V]): Tensor1[L, V] =
     def rescale(x: Tensor1[L, V]): Tensor1[L, V] =
-      // Unlike LayerNorm, RMSNorm does not re-center: it only divides by the root mean square.
       val meanSquare = x.pow(2).mean
       x /! (meanSquare + ε).sqrt
     rescale(x) * params.weight
diff --git a/core/src/main/scala/deepwit/optimizer/LearningRateSchedule.scala b/core/src/main/scala/deepwit/optimizer/LearningRateSchedule.scala
index b58f3cd..8fa2494 100644
--- a/core/src/main/scala/deepwit/optimizer/LearningRateSchedule.scala
+++ b/core/src/main/scala/deepwit/optimizer/LearningRateSchedule.scala
@@ -10,6 +10,10 @@ case class LearningRateSchedulerState[P, State[_]](
 )
 type LearningRateSchedulerStateFor[State[_]] = [P] =>> LearningRateSchedulerState[P, State]
 
+/** Wraps an optimizer so its learning rate follows `schedule`, counting steps from 1.
+  *
+  * @param optF Builds the optimizer at a given learning rate, e.g. `lr => Adam(lr)`.
+  */
 class LearningRateScheduler[State[_]](val optF: Tensor0[Float32] => GradientOptimizer[State], schedule: Tensor0[Int32] => Tensor0[Float32]) extends GradientOptimizer[LearningRateSchedulerStateFor[State]]:
 
   def init[P: TensorTree, V](params: P)(using TreeOf[P, V])(using IsFloating[V]): LearningRateSchedulerState[P, State] =
@@ -24,6 +28,11 @@ class LearningRateScheduler[State[_]](val optF: Tensor0[Float32] => GradientOpti
     val (newParams, newOptState) = opt.update(gradients, params, optState)
     (newParams, LearningRateSchedulerState(step + 1, newOptState))
 
+/** A learning rate as a function of the step count, which starts at 1.
+  *
+  * @param steps How many steps this schedule is defined for, after which it holds its final value.
+  *              [[LearningRateSchedule.followBy]] uses it to know when to hand over.
+  */
 trait LearningRateSchedule(val steps: Tensor0[Int32]) extends (Tensor0[Int32] => Tensor0[Float32])
 
 object LearningRateSchedule:
@@ -54,28 +63,21 @@ object LearningRateSchedule:
       *
       * For all `t < steps`, the schedule evaluates as if `t = 0`, effectively locking
       * the learning rate at its initial starting value until the delay has passed.
-      *
-      * @param delaySteps The number of iterations to delay the schedule's progression.
-      * @return A time-shifted schedule.
       */
     def delay(delaySteps: Tensor0[Int32]): LearningRateSchedule = DelayedSchedule(schedule, delaySteps)
 
+    /** Runs this schedule through its `steps`, then hands over to `second` delayed by that much,
+      * so `second` starts from its own beginning rather than mid-curve.
+      */
     def followBy(second: LearningRateSchedule): LearningRateSchedule = FollowBySchedule(schedule, second)
 
 object LearningRateSchedules:
 
-  /** Creates a schedule that maintains a constant learning rate for a specified number of steps. */
+  /** A schedule of constant `learningRate` for a specified number of `steps`. */
   class ConstantLearningRate(val learningRate: Tensor0[Float32], steps: Int = Int.MaxValue) extends LearningRateSchedule(steps = Tensor0(steps)):
     override def apply(step: Tensor0[Int32]): Tensor0[Float32] = learningRate
 
-  /** Creates a schedule that rises linearly from `minLr` to `maxLr` over a specified number of warmup steps.
-    *
-    * @param vtype The floating-point type to use for the learning rate values.
-    * @param from The initial learning rate at the start of the warmup (typically 0.0).
-    * @param to The peak learning rate reached at the end of the warmup.
-    * @param warmupSteps The number of steps over which the learning rate increases linearly.
-    * @return A linear warmup schedule.
-    */
+  /** A schedule rising linearly from `from` to `to` for a specified number of `warmupSteps`. */
   class LinearSchedule(
       val from: Tensor0[Float32],
       val to: Tensor0[Float32],
@@ -85,6 +87,8 @@ object LearningRateSchedules:
     private val vtype = from.vtype
 
     override def apply(step: Tensor0[Int32]): Tensor0[Float32] =
+      // Steps are 1-based, and the +1 keeps the first step off `from` — with LinearWarmup's
+      // `from = 0` that would be a step at zero learning rate.
       val warmupRatio = minimum((step.asFloat(vtype)) / ((warmupSteps + 1).asFloat(vtype)), 1f)
       from + warmupRatio * (to - from)
 
@@ -94,14 +98,9 @@ object LearningRateSchedules:
         warmupSteps: Tensor0[Int32]
     ): LinearSchedule = new LinearSchedule(0.0f, to, warmupSteps)
 
-  /** Creates a schedule that decays from `maxLr` down to `minLr` following a half-cosine curve.
-    *
-    * This schedule has no concept of warmup; it begins decaying immediately at `t = 0`.
-    * Once `t >= decaySteps`, the learning rate locks permanently at `minLr`.
+  /** A schedule that decays from `from` down to `to` for `decaySteps` following a half-cosine curve.
     *
-    * @param from The initial maximum learning rate at `t = 0`.
-    * @param to The final baseline learning rate to reach after decaying.
-    * @param decaySteps The number of steps over which to apply the decay curve.
+    * Past `decaySteps` the learning rate stays at `to`.
     */
   class CosineDecay(
       val from: Tensor0[Float32],
diff --git a/core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala b/core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala
new file mode 100644
index 0000000..e5910ec
--- /dev/null
+++ b/core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala
@@ -0,0 +1,56 @@
+package deepwit.attention
+
+import deepwit.*
+import dimwit.*
+import dimwit.stats.Normal
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.funspec.AnyFunSpec
+
+class ReferenceMultiHeadAttentionSuite extends AnyFunSpec with Matchers:
+
+  trait Src derives Label
+  trait SrcEmb derives Label
+  trait Tgt derives Label
+  trait TgtEmb derives Label
+
+  private val numHeads = 2
+  private val srcExtent = Axis[Src] -> 3
+  private val tgtExtent = Axis[Tgt] -> 5
+  private val srcEmbExtent = Axis[SrcEmb] -> 4
+  private val tgtEmbExtent = Axis[TgtEmb] -> 4
+
+  private val batched = MultiHeadAttention.Params.xavierUniformDepthScaled(
+    numTransformerLayers = 2,
+    numHeads = numHeads,
+    sourceEmbeddingExtent = srcEmbExtent,
+    targetEmbeddingExtent = tgtEmbExtent,
+    vtype = VType[Float32],
+    key = Random.Key(42)
+  )
+
+  /** The weights the batched implementation holds, split head by head, so both see the same model. */
+  private val perHead = ReferenceMultiHeadAttention.Params(
+    heads = (0 until numHeads).toList.map(head =>
+      Attention.Params(
+        batched.queryWeights.slice(Axis[Head].at(head)),
+        batched.keyWeights.slice(Axis[Head].at(head)),
+        batched.valueWeights.slice(Axis[Head].at(head))
+      )
+    ),
+    outputProjection = batched.outputProjection
+  )
+
+  private def source = Normal.standardNormal(Shape(srcExtent, srcEmbExtent)).sample(Random.Key(7))
+  private def target = Normal.standardNormal(Shape(tgtExtent, tgtEmbExtent)).sample(Random.Key(8))
+
+  describe("ReferenceMultiHeadAttention"):
+
+    it("computes what MultiHeadAttention computes, under a full mask"):
+      val folded = MultiHeadCustomAttention(batched, fullMask[Tgt, Src], AttentionScore.scaledDotProduct)
+      val reference = ReferenceMultiHeadAttention(perHead, fullMask[Tgt, Src], AttentionScore.scaledDotProduct)
+      reference(source, target) should approxEqual(folded(source, target), 1e-5f)
+
+    it("computes what MultiHeadAttention computes, under a causal mask"):
+      val folded = MultiHeadCustomAttention(batched, causalMask[Tgt, Src], AttentionScore.scaledDotProduct)
+      val reference = ReferenceMultiHeadAttention(perHead, causalMask[Tgt, Src], AttentionScore.scaledDotProduct)
+      reference(source, target) should approxEqual(folded(source, target), 1e-5f)