Prepare and release 0.1.0

.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
 

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. 

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
@@ -57,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

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],

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

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,

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,

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]):
 

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,

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,

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],

core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala

@@ -3,25 +3,27 @@ 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 Λ
 
+/** 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)) // 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))
 
+  /** 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)
 
 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] =

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)

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.")
+  /** 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 = threshold * (residual - 0.5f * threshold)
-    where(residual <= threshold, squared, absolute)
+    val absolute = transitionPoint * (residual - 0.5f * transitionPoint)
+    where(residual <= transitionPoint, squared, absolute)

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

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