effect-gpt 0.1.0

package/src/index.ts

@@ -0,0 +1,88 @@
+export { Vocab } from "./vocab/Vocab"
+export { tokenize } from "./tokenize/tokenize"
+export { Dataset, DatasetLoadError, DatasetParseError } from "./data/Dataset"
+export {
+  TrainingError,
+  TrainingDatasetError,
+  TrainingShapeError,
+  TrainingTokenizerError,
+  TrainingOptimizerError,
+  TrainingConfigError,
+  TrainingUnknownError
+} from "./errors"
+export * from "./config"
+
+export type { Tensor2D } from "./tensor/Tensor2D"
+export * as T2D from "./tensor/Tensor2D"
+export * as TensorOps from "./tensor/ops"
+export { ShapeError } from "./tensor/ops"
+export type { Rng } from "./tensor/random"
+export { seeded } from "./tensor/random"
+export { systemRng } from "./tensor/random"
+
+export type { ModelLayer } from "./model/ModelLayer"
+export { Embeddings } from "./model/Embeddings"
+export { SelfAttention } from "./model/SelfAttention"
+export { FeedForward } from "./model/FeedForward"
+export { LayerNorm } from "./model/LayerNorm"
+export { TransformerBlock } from "./model/TransformerBlock"
+export { OutputProjection } from "./model/OutputProjection"
+export { LLM } from "./model/LLM"
+
+export { Adam } from "./training/Adam"
+export { clipGlobalL2 } from "./training/clip"
+export { softmaxRows, crossEntropyLoss, dLogits } from "./training/loss"
+export {
+  train,
+  trainStream,
+  LLMService,
+  TrainingConfig,
+  makeLLMLayer,
+  makeTrainingConfigLayer
+} from "./training/train"
+
+export type {
+  LogLevel,
+  LoggerService,
+  LoggerServiceId,
+  RandomService,
+  RandomServiceId,
+  SeedService,
+  SeedServiceId,
+  MetricsService,
+  MetricsServiceId,
+  MetricsSnapshot,
+  Counter,
+  Gauge,
+  Histogram,
+  TimingResult
+} from "./services"
+export {
+  Logger,
+  ConsoleLoggerLive,
+  TerminalLoggerLive,
+  NullLoggerLive,
+  log,
+  debug,
+  info,
+  warn,
+  error,
+  Random,
+  SeededRandomLive,
+  SystemRandomLive,
+  next,
+  nextGaussian,
+  nextInt,
+  fork,
+  Seed,
+  SeedLayer,
+  useSeedRng,
+  Metrics,
+  InMemoryMetricsLive,
+  NoOpMetricsLive,
+  counter,
+  gauge,
+  histogram,
+  timed,
+  snapshot
+} from "./services"

package/src/model/Embeddings.ts

@@ -0,0 +1,108 @@
+import * as Effect from "effect/Effect"
+import * as FiberId from "effect/FiberId"
+import type { Tensor2D } from "../tensor/Tensor2D"
+import * as T from "../tensor/Tensor2D"
+import * as Ops from "../tensor/ops"
+import type { ShapeError } from "../tensor/ops"
+import type { ModelLayer } from "./ModelLayer"
+import { MAX_SEQ_LEN, EMBEDDING_DIM } from "../config"
+import { Adam } from "../training/Adam"
+import type { Rng } from "../tensor/random"
+
+export class Embeddings implements ModelLayer {
+  readonly _tag = "Embeddings"
+  tokenEmbeddings: Tensor2D
+  positionalEmbeddings: Tensor2D
+
+  private cache = new Map<number | string, Tensor2D>()
+  private lastCache: Tensor2D | null = null
+  tokenOptimizer: Adam
+  positionalOptimizer: Adam
+
+  constructor(vocabSize: number, embeddingDim: number = EMBEDDING_DIM, maxSeqLen: number = MAX_SEQ_LEN, rng: Rng) {
+    this.tokenEmbeddings = Ops.initNormal(vocabSize, embeddingDim, 0, 0.02, rng)
+    this.positionalEmbeddings = Ops.initNormal(maxSeqLen, embeddingDim, 0, 0.02, rng)
+    this.tokenOptimizer = Adam.make(vocabSize, embeddingDim)
+    this.positionalOptimizer = Adam.make(maxSeqLen, embeddingDim)
+  }
+
+  private fiberKey(fiberId: FiberId.FiberId): number | string {
+    return FiberId.isRuntime(fiberId) ? fiberId.id : JSON.stringify(fiberId)
+  }
+
+  get parametersCount(): number {
+    return this.tokenEmbeddings.data.length + this.positionalEmbeddings.data.length
+  }
+
+  forward(input: Tensor2D): Effect.Effect<Tensor2D, ShapeError> {
+    return Effect.gen(this, function* () {
+      const fiberId = yield* Effect.fiberId
+      const key = this.fiberKey(fiberId)
+      const cloned = T.clone(input)
+      this.cache.set(key, cloned)
+      this.lastCache = cloned
+      const tokenIds: Array<number> = []
+      for (let i = 0; i < input.data.length; i++) {
+        // Match Rust's float-to-usize truncation behavior.
+        tokenIds.push(Math.trunc(input.data[i]))
+      }
+
+      const seqLen = tokenIds.length
+      if (seqLen > this.positionalEmbeddings.rows) {
+        return yield* Effect.fail(
+          new Ops.ShapeError(`Sequence length ${seqLen} exceeds maximum ${this.positionalEmbeddings.rows}`)
+        )
+      }
+
+      const tokenEmbeds = yield* Ops.gatherRows(this.tokenEmbeddings, tokenIds)
+      const posEmbeds = yield* Ops.sliceRows(this.positionalEmbeddings, 0, seqLen)
+      const combined = yield* Ops.add(tokenEmbeds, posEmbeds)
+      return combined
+    })
+  }
+
+  backward(dOut: Tensor2D, lr: number): Effect.Effect<Tensor2D, ShapeError> {
+    return Effect.gen(this, function* () {
+      const fiberId = yield* Effect.fiberId
+      const key = this.fiberKey(fiberId)
+      const cachedInput = this.cache.get(key) ?? this.lastCache
+      if (!cachedInput) {
+        return yield* Effect.fail(new Ops.ShapeError("Embeddings.backward called before forward"))
+      }
+      this.cache.delete(key)
+      this.lastCache = null
+
+      const input = cachedInput
+      const tokenIds: Array<number> = []
+      for (let i = 0; i < input.data.length; i++) {
+        tokenIds.push(Math.trunc(input.data[i]))
+      }
+
+      const tokenGrads = T.zeros(this.tokenEmbeddings.rows, this.tokenEmbeddings.cols)
+      const positionalGrads = T.zeros(this.positionalEmbeddings.rows, this.positionalEmbeddings.cols)
+
+      const seqLen = tokenIds.length
+      for (let i = 0; i < seqLen; i++) {
+        const tokenId = tokenIds[i]
+        if (tokenId < 0 || tokenId >= this.tokenEmbeddings.rows) {
+          return yield* Effect.fail(
+            new Ops.ShapeError(`Token ID ${tokenId} out of bounds for vocab size ${this.tokenEmbeddings.rows}`)
+          )
+        }
+        const rowOffset = i * dOut.cols
+        const tokenOffset = tokenId * tokenGrads.cols
+        const posOffset = i * positionalGrads.cols
+        for (let j = 0; j < dOut.cols; j++) {
+          const grad = dOut.data[rowOffset + j]
+          tokenGrads.data[tokenOffset + j] += grad
+          positionalGrads.data[posOffset + j] += grad
+        }
+      }
+
+      this.tokenOptimizer.step(this.tokenEmbeddings, tokenGrads, lr)
+      this.positionalOptimizer.step(this.positionalEmbeddings, positionalGrads, lr)
+
+      return T.clone(dOut)
+    })
+  }
+}

package/src/model/FeedForward.ts

@@ -0,0 +1,121 @@
+import * as Effect from "effect/Effect"
+import * as FiberId from "effect/FiberId"
+import type { Tensor2D } from "../tensor/Tensor2D"
+import * as T from "../tensor/Tensor2D"
+import * as Ops from "../tensor/ops"
+import type { ShapeError } from "../tensor/ops"
+import type { ModelLayer } from "./ModelLayer"
+import { EMBEDDING_DIM, HIDDEN_DIM } from "../config"
+import { Adam } from "../training/Adam"
+import type { Rng } from "../tensor/random"
+
+export class FeedForward implements ModelLayer {
+  readonly _tag = "FeedForward"
+  w1: Tensor2D
+  b1: Tensor2D
+  w2: Tensor2D
+  b2: Tensor2D
+
+  private cache = new Map<
+    number | string,
+    { input: Tensor2D; hiddenPreActivation: Tensor2D; hiddenPostActivation: Tensor2D }
+  >()
+  private lastCache:
+    | { input: Tensor2D; hiddenPreActivation: Tensor2D; hiddenPostActivation: Tensor2D }
+    | null = null
+  optimizerW1: Adam
+  optimizerB1: Adam
+  optimizerW2: Adam
+  optimizerB2: Adam
+
+  constructor(embeddingDim: number = EMBEDDING_DIM, hiddenDim: number = HIDDEN_DIM, rng: Rng) {
+    const stdW1 = Math.sqrt(2.0 / embeddingDim)
+    const stdW2 = Math.sqrt(2.0 / hiddenDim)
+
+    this.w1 = Ops.initNormal(embeddingDim, hiddenDim, 0, stdW1, rng)
+    this.b1 = T.zeros(1, hiddenDim)
+    this.w2 = Ops.initNormal(hiddenDim, embeddingDim, 0, stdW2, rng)
+    this.b2 = T.zeros(1, embeddingDim)
+    this.optimizerW1 = Adam.make(embeddingDim, hiddenDim)
+    this.optimizerB1 = Adam.make(1, hiddenDim)
+    this.optimizerW2 = Adam.make(hiddenDim, embeddingDim)
+    this.optimizerB2 = Adam.make(1, embeddingDim)
+  }
+
+  private fiberKey(fiberId: FiberId.FiberId): number | string {
+    return FiberId.isRuntime(fiberId) ? fiberId.id : JSON.stringify(fiberId)
+  }
+
+  get parametersCount(): number {
+    return this.w1.data.length + this.b1.data.length + this.w2.data.length + this.b2.data.length
+  }
+
+  forward(input: Tensor2D): Effect.Effect<Tensor2D, ShapeError> {
+    return Effect.gen(this, function* () {
+      const fiberId = yield* Effect.fiberId
+      const key = this.fiberKey(fiberId)
+
+      const h1 = yield* Ops.matMul(input, this.w1)
+      const h1Bias = yield* Ops.addRowBias(h1, this.b1)
+      const h1BiasClone = T.clone(h1Bias)
+
+      const h1Relu = Ops.relu(h1Bias)
+      const h1ReluClone = T.clone(h1Relu)
+      const cached = {
+        input: T.clone(input),
+        hiddenPreActivation: h1BiasClone,
+        hiddenPostActivation: h1ReluClone
+      }
+      this.cache.set(key, cached)
+      this.lastCache = cached
+
+      const h2 = yield* Ops.matMul(h1Relu, this.w2)
+      const h2Bias = yield* Ops.addRowBias(h2, this.b2)
+      const output = yield* Ops.add(h2Bias, input)
+      return output
+    })
+  }
+
+  backward(dOut: Tensor2D, lr: number): Effect.Effect<Tensor2D, ShapeError> {
+    return Effect.gen(this, function* () {
+      const fiberId = yield* Effect.fiberId
+      const key = this.fiberKey(fiberId)
+      const cached = this.cache.get(key) ?? this.lastCache
+      if (!cached) {
+        return yield* Effect.fail(new Ops.ShapeError("FeedForward.backward called before forward"))
+      }
+      this.cache.delete(key)
+      this.lastCache = null
+
+      const { input, hiddenPreActivation, hiddenPostActivation } = cached
+
+      const hiddenPostT = Ops.transpose(hiddenPostActivation)
+      const gradW2 = yield* Ops.matMul(hiddenPostT, dOut)
+      const gradB2 = Ops.sumCols(dOut)
+
+      const w2T = Ops.transpose(this.w2)
+      const gradHiddenPost = yield* Ops.matMul(dOut, w2T)
+
+      const reluGrad = T.zeros(hiddenPreActivation.rows, hiddenPreActivation.cols)
+      for (let i = 0; i < hiddenPreActivation.data.length; i++) {
+        reluGrad.data[i] = hiddenPreActivation.data[i] > 0 ? 1 : 0
+      }
+      const gradHiddenPre = yield* Ops.mul(gradHiddenPost, reluGrad)
+
+      const inputT = Ops.transpose(input)
+      const gradW1 = yield* Ops.matMul(inputT, gradHiddenPre)
+      const gradB1 = Ops.sumCols(gradHiddenPre)
+
+      const w1T = Ops.transpose(this.w1)
+      const gradInputFF = yield* Ops.matMul(gradHiddenPre, w1T)
+      const gradInput = yield* Ops.add(gradInputFF, dOut)
+
+      this.optimizerW2.step(this.w2, gradW2, lr)
+      this.optimizerB2.step(this.b2, gradB2, lr)
+      this.optimizerW1.step(this.w1, gradW1, lr)
+      this.optimizerB1.step(this.b1, gradB1, lr)
+
+      return gradInput
+    })
+  }
+}

package/src/model/LLM.ts

@@ -0,0 +1,124 @@
+import * as Effect from "effect/Effect"
+import * as Option from "effect/Option"
+import type { Tensor2D } from "../tensor/Tensor2D"
+import * as T from "../tensor/Tensor2D"
+import * as Ops from "../tensor/ops"
+import type { ShapeError } from "../tensor/ops"
+import type { ModelLayer } from "./ModelLayer"
+import { Embeddings } from "./Embeddings"
+import { TransformerBlock } from "./TransformerBlock"
+import { OutputProjection } from "./OutputProjection"
+import { Vocab } from "../vocab/Vocab"
+import { tokenize } from "../tokenize/tokenize"
+import { MAX_SEQ_LEN, EMBEDDING_DIM, HIDDEN_DIM } from "../config"
+import type { Rng } from "../tensor/random"
+
+export class LLM {
+  readonly vocab: Vocab
+  readonly network: ReadonlyArray<ModelLayer>
+
+  constructor(vocab: Vocab, network: ReadonlyArray<ModelLayer>) {
+    this.vocab = vocab
+    this.network = network
+  }
+
+  static default(rng: Rng, numTransformerBlocks = 1): LLM {
+    const vocab = Vocab.make(Vocab.defaultWords())
+    return LLM.make(vocab, rng, numTransformerBlocks)
+  }
+
+  static make(vocab: Vocab, rng: Rng, numTransformerBlocks = 3): LLM {
+    const vocabSize = vocab.words.length
+    const network: Array<ModelLayer> = [
+      new Embeddings(vocabSize, EMBEDDING_DIM, MAX_SEQ_LEN, rng),
+      ...Array.from({ length: numTransformerBlocks }, () => new TransformerBlock(EMBEDDING_DIM, HIDDEN_DIM, rng)),
+      new OutputProjection(EMBEDDING_DIM, vocabSize, rng)
+    ]
+    return new LLM(vocab, network)
+  }
+
+  networkDescription(): string {
+    return this.network.map((layer) => layer._tag).join(", ")
+  }
+
+  totalParameters(): number {
+    return this.network.reduce((sum, layer) => sum + layer.parametersCount, 0)
+  }
+
+  predict(text: string): Effect.Effect<string, ShapeError> {
+    return Effect.gen(this, function* () {
+      const outputTokens = yield* this.forward(text)
+
+      if (outputTokens.length === 0) {
+        return ""
+      }
+
+      const tokenStrs: Array<string> = []
+      for (const t of outputTokens) {
+        const decoded = this.vocab.decode(t)
+        if (Option.isSome(decoded)) {
+          tokenStrs.push(decoded.value)
+        }
+      }
+
+      return tokenStrs.join(" ")
+    })
+  }
+
+  forward(text: string): Effect.Effect<ReadonlyArray<number>, ShapeError> {
+    return Effect.gen(this, function* () {
+      const tokenized: Array<number> = [...tokenize(text, this.vocab)]
+      const outputTokens: Array<number> = []
+
+      if (tokenized.length === 0) {
+        return outputTokens
+      }
+
+      const inputLen = tokenized.length
+      if (inputLen >= MAX_SEQ_LEN) {
+        return outputTokens
+      }
+
+      const endTokenId = this.vocab.encode("</s>")
+      if (Option.isNone(endTokenId)) {
+        return yield* Effect.fail(new Ops.ShapeError("End token </s> not found in vocabulary"))
+      }
+
+      for (let step = 0; step < MAX_SEQ_LEN - inputLen; step++) {
+        if (outputTokens.length >= MAX_SEQ_LEN - 1) {
+          break
+        }
+
+        const tokenInput = T.fromArray(1, tokenized.length, tokenized)
+        let input: Tensor2D = tokenInput
+
+        for (const layer of this.network) {
+          input = yield* layer.forward(input)
+        }
+
+        const logits = input
+
+        if (logits.rows === 0) {
+          break
+        }
+
+        const lastRowStart = (logits.rows - 1) * logits.cols
+        const lastLogitData = logits.data.slice(lastRowStart, lastRowStart + logits.cols)
+        const lastLogit = T.make(1, logits.cols, lastLogitData)
+
+        const probs = Ops.softmaxRows(lastLogit)
+        const tokens = Ops.argmaxRows(probs)
+        const nextToken = tokens[tokens.length - 1]
+
+        outputTokens.push(nextToken)
+        tokenized.push(nextToken)
+
+        if (nextToken === endTokenId.value) {
+          break
+        }
+      }
+
+      return outputTokens
+    })
+  }
+}

package/src/model/LayerNorm.ts

@@ -0,0 +1,138 @@
+import * as Effect from "effect/Effect"
+import * as FiberId from "effect/FiberId"
+import type { Tensor2D } from "../tensor/Tensor2D"
+import * as T from "../tensor/Tensor2D"
+import * as Ops from "../tensor/ops"
+import type { ShapeError } from "../tensor/ops"
+import type { ModelLayer } from "./ModelLayer"
+import { Adam } from "../training/Adam"
+
+export class LayerNorm implements ModelLayer {
+  readonly _tag = "LayerNorm"
+  readonly epsilon: number = 1e-5
+  gamma: Tensor2D
+  beta: Tensor2D
+
+  private cache = new Map<number | string, { input: Tensor2D; mean: Tensor2D; variance: Tensor2D }>()
+  private lastCache: { input: Tensor2D; mean: Tensor2D; variance: Tensor2D } | null = null
+  optimizerGamma: Adam
+  optimizerBeta: Adam
+
+  constructor(embeddingDim: number) {
+    this.gamma = T.ones(1, embeddingDim)
+    this.beta = T.zeros(1, embeddingDim)
+    this.optimizerGamma = Adam.make(1, embeddingDim)
+    this.optimizerBeta = Adam.make(1, embeddingDim)
+  }
+
+  private fiberKey(fiberId: FiberId.FiberId): number | string {
+    return FiberId.isRuntime(fiberId) ? fiberId.id : JSON.stringify(fiberId)
+  }
+
+  get parametersCount(): number {
+    return this.gamma.data.length + this.beta.data.length
+  }
+
+  forward(input: Tensor2D): Effect.Effect<Tensor2D, ShapeError> {
+    return Effect.gen(this, function* () {
+      const mean = Ops.meanRows(input)
+      const variance = Ops.varRows(input)
+
+      const fiberId = yield* Effect.fiberId
+      const key = this.fiberKey(fiberId)
+      const cached = {
+        input: T.clone(input),
+        mean: T.clone(mean),
+        variance: T.clone(variance)
+      }
+      this.cache.set(key, cached)
+      this.lastCache = cached
+
+      // Use sqrt(variance + epsilon) for numerical stability
+      const rstd = Ops.mapScalar(variance, (v) => 1.0 / Math.sqrt(v + this.epsilon))
+      const centered = yield* Ops.broadcastSubCol(input, mean)
+      const normalized = yield* Ops.broadcastMulCol(centered, rstd)
+      const scaled = yield* Ops.broadcastMulRow(normalized, this.gamma)
+      const shifted = yield* Ops.broadcastAddRow(scaled, this.beta)
+      return shifted
+    })
+  }
+
+  backward(dOut: Tensor2D, lr: number): Effect.Effect<Tensor2D, ShapeError> {
+    return Effect.gen(this, function* () {
+      const fiberId = yield* Effect.fiberId
+      const key = this.fiberKey(fiberId)
+      const cached = this.cache.get(key) ?? this.lastCache
+      if (!cached) {
+        return yield* Effect.fail(new Ops.ShapeError("LayerNorm.backward called before forward"))
+      }
+      this.cache.delete(key)
+      this.lastCache = null
+
+      const { input, mean, variance } = cached
+      const rows = input.rows
+      const cols = input.cols
+      const nFeatures = cols
+
+      const normalized = T.zeros(rows, cols)
+      const gradNormalized = T.zeros(rows, cols)
+      for (let i = 0; i < rows; i++) {
+        const meanVal = mean.data[i]
+        // Consistent with forward: rstd = 1 / sqrt(variance + epsilon)
+        const rstd = 1.0 / Math.sqrt(variance.data[i] + this.epsilon)
+        for (let j = 0; j < cols; j++) {
+          const idx = i * cols + j
+          const norm = (input.data[idx] - meanVal) * rstd
+          normalized.data[idx] = norm
+          gradNormalized.data[idx] = this.gamma.data[j] * dOut.data[idx]
+        }
+      }
+
+      const gradGamma = T.zeros(1, cols)
+      const gradBeta = T.zeros(1, cols)
+      for (let j = 0; j < cols; j++) {
+        let sumGamma = 0
+        let sumBeta = 0
+        for (let i = 0; i < rows; i++) {
+          const idx = i * cols + j
+          sumGamma += normalized.data[idx] * dOut.data[idx]
+          sumBeta += dOut.data[idx]
+        }
+        gradGamma.data[j] = sumGamma
+        gradBeta.data[j] = sumBeta
+      }
+
+      const gradInput = T.zeros(rows, cols)
+      for (let i = 0; i < rows; i++) {
+        const meanVal = mean.data[i]
+        const varPlusEps = variance.data[i] + this.epsilon
+        const rstd = 1.0 / Math.sqrt(varPlusEps)
+
+        let sumGradNormalized = 0
+        let sumGradNormTimesNorm = 0
+        const rowOffset = i * cols
+
+        for (let j = 0; j < cols; j++) {
+          const idx = rowOffset + j
+          sumGradNormalized += gradNormalized.data[idx]
+          sumGradNormTimesNorm += gradNormalized.data[idx] * normalized.data[idx]
+        }
+
+        // Gradient of LayerNorm: dL/dx = rstd * (dL/dnorm - mean(dL/dnorm) - norm * mean(dL/dnorm * norm))
+        for (let j = 0; j < cols; j++) {
+          const idx = rowOffset + j
+          gradInput.data[idx] =
+            rstd *
+            (gradNormalized.data[idx] -
+              sumGradNormalized / nFeatures -
+              (normalized.data[idx] * sumGradNormTimesNorm) / nFeatures)
+        }
+      }
+
+      this.optimizerGamma.step(this.gamma, gradGamma, lr)
+      this.optimizerBeta.step(this.beta, gradBeta, lr)
+
+      return gradInput
+    })
+  }
+}

package/src/model/ModelLayer.ts

@@ -0,0 +1,10 @@
+import type * as Effect from "effect/Effect"
+import type { Tensor2D } from "../tensor/Tensor2D"
+import type { ShapeError } from "../tensor/ops"
+
+export interface ModelLayer {
+  readonly _tag: string
+  readonly parametersCount: number
+  forward(input: Tensor2D): Effect.Effect<Tensor2D, ShapeError>
+  backward(dOut: Tensor2D, lr: number): Effect.Effect<Tensor2D, ShapeError>
+}

package/src/model/OutputProjection.ts

@@ -0,0 +1,76 @@
+import * as Effect from "effect/Effect"
+import * as FiberId from "effect/FiberId"
+import type { Tensor2D } from "../tensor/Tensor2D"
+import * as T from "../tensor/Tensor2D"
+import * as Ops from "../tensor/ops"
+import type { ShapeError } from "../tensor/ops"
+import type { ModelLayer } from "./ModelLayer"
+import { EMBEDDING_DIM } from "../config"
+import { Adam } from "../training/Adam"
+import type { Rng } from "../tensor/random"
+
+export class OutputProjection implements ModelLayer {
+  readonly _tag = "OutputProjection"
+  wOut: Tensor2D
+  bOut: Tensor2D
+
+  private cache = new Map<number | string, Tensor2D>()
+  private lastCache: Tensor2D | null = null
+  optimizerWOut: Adam
+
+  constructor(embeddingDim: number = EMBEDDING_DIM, vocabSize: number, rng: Rng) {
+    const std = Math.sqrt(2.0 / embeddingDim)
+    this.wOut = Ops.initNormal(embeddingDim, vocabSize, 0, std, rng)
+    this.bOut = T.zeros(1, vocabSize)
+    this.optimizerWOut = Adam.make(embeddingDim, vocabSize)
+  }
+
+  private fiberKey(fiberId: FiberId.FiberId): number | string {
+    return FiberId.isRuntime(fiberId) ? fiberId.id : JSON.stringify(fiberId)
+  }
+
+  get parametersCount(): number {
+    return this.wOut.data.length + this.bOut.data.length
+  }
+
+  forward(input: Tensor2D): Effect.Effect<Tensor2D, ShapeError> {
+    return Effect.gen(this, function* () {
+      const fiberId = yield* Effect.fiberId
+      const key = this.fiberKey(fiberId)
+      const cloned = T.clone(input)
+      this.cache.set(key, cloned)
+      this.lastCache = cloned
+      const projected = yield* Ops.matMul(input, this.wOut)
+      const output = yield* Ops.addRowBias(projected, this.bOut)
+      return output
+    })
+  }
+
+  backward(dOut: Tensor2D, lr: number): Effect.Effect<Tensor2D, ShapeError> {
+    return Effect.gen(this, function* () {
+      const fiberId = yield* Effect.fiberId
+      const key = this.fiberKey(fiberId)
+      const cachedInput = this.cache.get(key) ?? this.lastCache
+      if (!cachedInput) {
+        return yield* Effect.fail(new Ops.ShapeError("OutputProjection.backward called before forward"))
+      }
+      this.cache.delete(key)
+      this.lastCache = null
+
+      const input = cachedInput
+      const inputT = Ops.transpose(input)
+      const gradWOut = yield* Ops.matMul(inputT, dOut)
+      const gradBOut = Ops.sumCols(dOut)
+
+      const wOutT = Ops.transpose(this.wOut)
+      const gradInput = yield* Ops.matMul(dOut, wOutT)
+
+      this.optimizerWOut.step(this.wOut, gradWOut, lr)
+      for (let j = 0; j < this.bOut.data.length; j++) {
+        this.bOut.data[j] -= lr * gradBOut.data[j]
+      }
+
+      return gradInput
+    })
+  }
+}