effect-gpt 0.1.0

package/src/training/train.ts

@@ -0,0 +1,203 @@
+import * as Effect from "effect/Effect"
+import * as Stream from "effect/Stream"
+import * as Chunk from "effect/Chunk"
+import * as Context from "effect/Context"
+import * as Layer from "effect/Layer"
+import * as Ref from "effect/Ref"
+import * as Option from "effect/Option"
+import type { ShapeError } from "../tensor/ops"
+import * as Ops from "../tensor/ops"
+import * as T from "../tensor/Tensor2D"
+import { tokenize } from "../tokenize/tokenize"
+import type { LLM } from "../model/LLM"
+import { softmaxRows, crossEntropyLoss, dLogits } from "./loss"
+import { clipGlobalL2 } from "./clip"
+import type { LoggerServiceId } from "../services/Logger"
+import { info } from "../services/Logger"
+import type { MetricsServiceId } from "../services/Metrics"
+import { counter, gauge, timed } from "../services/Metrics"
+import { TrainingError } from "../errors"
+import type { TrainingError as TrainingErrorType } from "../errors"
+
+export interface TrainingConfig {
+  readonly epochs: number
+  readonly learningRate: number
+  readonly clipNorm?: number
+  readonly preprocessConcurrency?: number | "unbounded"
+  readonly preprocessBatchSize?: number
+  readonly trainConcurrency?: number
+}
+
+interface TrainingConfigId {
+  readonly TrainingConfig: unique symbol
+}
+interface LLMServiceId {
+  readonly LLMService: unique symbol
+}
+
+export interface PreprocessSettings {
+  readonly concurrency: number | "unbounded"
+  readonly batchSize: number
+}
+
+interface PreprocessSettingsId {
+  readonly PreprocessSettings: unique symbol
+}
+
+export const TrainingConfig = Context.GenericTag<TrainingConfigId, TrainingConfig>("TrainingConfig")
+export const LLMService = Context.GenericTag<LLMServiceId, LLM>("LLMService")
+export const PreprocessSettings = Context.GenericTag<PreprocessSettingsId, PreprocessSettings>(
+  "PreprocessSettings"
+)
+
+export const makeLLMLayer = (llm: LLM) => Layer.succeed(LLMService, llm)
+export const makeTrainingConfigLayer = (config: TrainingConfig) =>
+  Layer.succeed(TrainingConfig, config)
+export const makePreprocessSettingsLayer = (settings: PreprocessSettings) =>
+  Layer.succeed(PreprocessSettings, settings)
+
+type TrainEnv =
+  | TrainingConfigId
+  | LLMServiceId
+  | LoggerServiceId
+  | MetricsServiceId
+  | PreprocessSettingsId
+
+const mapShapeError = <A, R>(effect: Effect.Effect<A, ShapeError, R>) =>
+  effect.pipe(Effect.mapError(TrainingError.shape))
+
+const mapShapeUnknown = (error: unknown): TrainingErrorType =>
+  error instanceof Ops.ShapeError ? TrainingError.shape(error) : TrainingError.fromUnknown(error)
+
+const wrapThrowing = <A>(
+  thunk: () => A,
+  mapError: (error: unknown) => TrainingErrorType = TrainingError.fromUnknown
+) =>
+  Effect.try({
+    try: thunk,
+    catch: (error) => mapError(error)
+  })
+
+const clampConcurrency = (value: number | undefined, fallback: number): number =>
+  value === undefined ? fallback : Math.max(1, value)
+
+const trainWithStreamFactory = <E, R>(
+  makeStream: () => Stream.Stream<string, E, R>
+): Effect.Effect<void, TrainingErrorType, R | TrainEnv> =>
+  Effect.gen(function* () {
+    const llm = yield* LLMService
+    const config = yield* TrainingConfig
+    const preprocessSettings = yield* Effect.gen(function* () {
+      const env = (yield* Effect.context<R | TrainEnv>()) as Context.Context<R | TrainEnv>
+      const maybeSettings = Context.getOption(env, PreprocessSettings)
+      if (maybeSettings._tag === "Some") {
+        return maybeSettings.value
+      }
+      return {
+        concurrency: config.preprocessConcurrency ?? "unbounded",
+        batchSize: config.preprocessBatchSize ?? 1
+      } satisfies PreprocessSettings
+    })
+
+    const endTokenId = llm.vocab.encode("</s>")
+    if (endTokenId._tag === "None") {
+      return yield* Effect.fail(TrainingError.config("End token </s> not found in vocabulary"))
+    }
+
+    const clipNorm = config.clipNorm ?? 5.0
+    const concurrency = preprocessSettings.concurrency
+    const batchSize = Math.max(1, preprocessSettings.batchSize)
+    const trainConcurrency = clampConcurrency(config.trainConcurrency, 4)
+
+    const epochCounter = yield* counter("epochs_completed")
+    const lossGauge = yield* gauge("epoch_loss")
+    const examplesCounter = yield* counter("examples_processed")
+
+    for (let epoch = 0; epoch < config.epochs; epoch++) {
+      const epochResult = yield* timed(`epoch_${epoch}`, Effect.gen(function* () {
+        const totalLossRef = yield* Ref.make(0)
+        const totalExamplesRef = yield* Ref.make(0)
+
+        const preprocess = (text: string) =>
+          Effect.sync(() => {
+            const tokens = [...tokenize(text, llm.vocab)]
+            if (tokens.length < 2) {
+              return Option.none<{ inputIds: number[]; targetIds: number[] }>()
+            }
+
+            return Option.some({
+              inputIds: tokens.slice(0, tokens.length - 1),
+              targetIds: tokens.slice(1)
+            })
+        })
+
+        const preprocessed = makeStream()
+          .pipe(
+            Stream.mapError(TrainingError.fromUnknown),
+            Stream.mapChunks(Chunk.chunksOf(batchSize)),
+            Stream.flattenChunks,
+            Stream.mapEffect(preprocess, { concurrency }),
+            Stream.filterMap((value) => value)
+          )
+
+        const trainExample = ({ inputIds, targetIds }: { inputIds: number[]; targetIds: number[] }) =>
+          Effect.gen(function* () {
+            let input = T.fromArray(1, inputIds.length, inputIds)
+            for (const layer of llm.network) {
+              input = yield* mapShapeError(layer.forward(input))
+            }
+
+            const logits = input
+            const probs = yield* wrapThrowing(() => softmaxRows(logits), mapShapeUnknown)
+            const loss = yield* wrapThrowing(() => crossEntropyLoss(probs, targetIds), mapShapeUnknown)
+            yield* Ref.update(totalLossRef, (current) => current + loss)
+            yield* Ref.update(totalExamplesRef, (current) => current + 1)
+
+            let grads = yield* wrapThrowing(() => dLogits(probs, targetIds), mapShapeUnknown)
+            clipGlobalL2(grads, clipNorm)
+
+            for (let i = llm.network.length - 1; i >= 0; i--) {
+              grads = yield* mapShapeError(llm.network[i]!.backward(grads, config.learningRate))
+            }
+
+            const tokens = Ops.argmaxRows(probs)
+            const nextToken = tokens[tokens.length - 1]
+            if (nextToken === endTokenId.value) {
+              return
+            }
+          })
+
+        yield* Effect.scoped(
+          Stream.runDrain(
+            Stream.mapEffect(preprocessed, trainExample, { concurrency: trainConcurrency })
+          )
+        )
+
+        const totalLoss = yield* Ref.get(totalLossRef)
+        const totalExamples = yield* Ref.get(totalExamplesRef)
+        yield* examplesCounter.inc(totalExamples)
+        return { totalLoss, totalExamples }
+      }))
+
+      const { totalLoss, totalExamples } = epochResult.value
+      const avgLoss = totalExamples > 0 ? totalLoss / totalExamples : 0
+
+      yield* lossGauge.set(avgLoss)
+      yield* epochCounter.inc()
+      yield* info(`Epoch ${epoch}: Loss = ${avgLoss.toFixed(4)}`, {
+        epoch,
+        loss: avgLoss,
+        examples: totalExamples,
+        durationMs: epochResult.durationMs
+      })
+    }
+  })
+
+export const train = (
+  examples: ReadonlyArray<string>
+): Effect.Effect<void, TrainingErrorType, TrainEnv> =>
+  trainWithStreamFactory(() => Stream.fromIterable(examples))
+
+export const trainStream = <E, R>(
+  makeStream: () => Stream.Stream<string, E, R>
+): Effect.Effect<void, TrainingErrorType, R | TrainEnv> => trainWithStreamFactory(makeStream)

package/src/vocab/Vocab.ts

@@ -0,0 +1,79 @@
+import * as HashMap from "effect/HashMap"
+import * as HashSet from "effect/HashSet"
+import * as Option from "effect/Option"
+import * as Stream from "effect/Stream"
+import * as Effect from "effect/Effect"
+import { splitWordToTokens } from "../tokenize/split"
+
+export class Vocab {
+  readonly encodeMap: HashMap.HashMap<string, number>
+  readonly decodeMap: HashMap.HashMap<number, string>
+  readonly words: ReadonlyArray<string>
+
+  private constructor(
+    encodeMap: HashMap.HashMap<string, number>,
+    decodeMap: HashMap.HashMap<number, string>,
+    words: ReadonlyArray<string>
+  ) {
+    this.encodeMap = encodeMap
+    this.decodeMap = decodeMap
+    this.words = words
+  }
+
+  static make(words: ReadonlyArray<string>): Vocab {
+    let encodeMap = HashMap.empty<string, number>()
+    let decodeMap = HashMap.empty<number, string>()
+
+    for (let i = 0; i < words.length; i += 1) {
+      const word = words[i]
+      encodeMap = HashMap.set(encodeMap, word, i)
+      decodeMap = HashMap.set(decodeMap, i, word)
+    }
+
+    return new Vocab(encodeMap, decodeMap, words)
+  }
+
+  static defaultWords(): ReadonlyArray<string> {
+    return ["hello", "world", "this", "is", "rust", "</s>"]
+  }
+
+  encode(word: string): Option.Option<number> {
+    return HashMap.get(this.encodeMap, word)
+  }
+
+  decode(id: number): Option.Option<string> {
+    return HashMap.get(this.decodeMap, id)
+  }
+
+  private static addTokensToSet(
+    set: HashSet.HashSet<string>,
+    text: string
+  ): HashSet.HashSet<string> {
+    const words = text.split(/\s+/).filter((w) => w.length > 0)
+    let next = set
+
+    for (const word of words) {
+      const parts = splitWordToTokens(word)
+      for (const part of parts) {
+        next = HashSet.add(next, part)
+      }
+    }
+
+    return next
+  }
+
+  static processTextForVocab(texts: ReadonlyArray<string>): HashSet.HashSet<string> {
+    let vocabSet = HashSet.add(HashSet.empty<string>(), "</s>")
+    for (const text of texts) {
+      vocabSet = Vocab.addTokensToSet(vocabSet, text)
+    }
+    return vocabSet
+  }
+
+  static processStreamForVocab<E, R>(
+    stream: Stream.Stream<string, E, R>
+  ): Effect.Effect<HashSet.HashSet<string>, E, R> {
+    const initial = HashSet.add(HashSet.empty<string>(), "</s>")
+    return Stream.runFold(stream, initial, (set, text) => Vocab.addTokensToSet(set, text))
+  }
+}

package/tests/fixtures/csv_bad.csv

@@ -0,0 +1,2 @@
+hello
+"unclosed

package/tests/fixtures/csv_good.csv

@@ -0,0 +1,3 @@
+hello,world
+"quoted,field",foo
+"say ""hi""",bar

package/tests/ts/cli_error_format.test.ts

@@ -0,0 +1,26 @@
+import { describe, expect, test } from "bun:test"
+import { formatTrainingError } from "../../src/cli/errors"
+import { TrainingError } from "../../src/errors"
+import { ShapeError } from "../../src/tensor/ops"
+
+describe("formatTrainingError", () => {
+  test("formats dataset error with path", () => {
+    const err = TrainingError.dataset({ _tag: "DatasetLoadError", path: "/tmp/data", error: new Error("io") } as any)
+    const message = formatTrainingError(err)
+    expect(message).toContain("Dataset error")
+    expect(message).toContain("/tmp/data")
+  })
+
+  test("formats shape error cause", () => {
+    const err = TrainingError.shape(new ShapeError("bad shape"))
+    const message = formatTrainingError(err)
+    expect(message).toContain("Shape error")
+    expect(message).toContain("bad shape")
+  })
+
+  test("formats plain Error as unexpected error", () => {
+    const message = formatTrainingError(new Error("boom"))
+    expect(message).toContain("Unexpected error")
+    expect(message).toContain("boom")
+  })
+})

package/tests/ts/dataset.test.ts

@@ -0,0 +1,35 @@
+import { describe, test, expect } from "bun:test"
+import * as Effect from "effect/Effect"
+import * as Exit from "effect/Exit"
+import { BunFileSystem } from "@effect/platform-bun"
+import { Dataset, DatasetParseError } from "../../src/data/Dataset.js"
+
+describe("Dataset CSV", () => {
+  const good = "tests/fixtures/csv_good.csv"
+  const bad = "tests/fixtures/csv_bad.csv"
+
+  test("collect reads CSV and joins fields", async () => {
+    const streams = Dataset.load({ pretrainingPath: good, chatPath: good, format: "csv" })
+    const program = Dataset.collect(streams).pipe(Effect.provide(BunFileSystem.layer))
+    const result = await Effect.runPromise(program)
+
+    expect(result.pretrainingData).toEqual(["hello,world", "quoted,field,foo", 'say "hi",bar'])
+    expect(result.chatTrainingData).toEqual(["hello,world", "quoted,field,foo", 'say "hi",bar'])
+  })
+
+  test("malformed CSV surfaces DatasetParseError", async () => {
+    const streams = Dataset.load({ pretrainingPath: bad, chatPath: bad, format: "csv" })
+    const program = Dataset.collect(streams).pipe(Effect.provide(BunFileSystem.layer))
+    const exit = await Effect.runPromiseExit(program)
+
+    expect(exit._tag).toBe("Failure")
+    if (exit._tag === "Failure") {
+      const cause = exit.cause
+      if (cause._tag === "Fail") {
+        expect(cause.error).toBeInstanceOf(DatasetParseError)
+      } else {
+        throw new Error(`Unexpected failure cause: ${cause._tag}`)
+      }
+    }
+  })
+})

package/tests/ts/embeddings.test.ts

@@ -0,0 +1,81 @@
+import { describe, test, expect } from "bun:test"
+import { runEffect } from "./support/runEffect"
+import { expectShape, expectNotClose, expectFinite } from "./support/tensorMatchers"
+import { makeEmbeddings } from "./support/factories"
+import * as T from "../../src/tensor/Tensor2D"
+import { EMBEDDING_DIM, MAX_SEQ_LEN } from "../../src/config"
+
+describe("Embeddings", () => {
+  test("embed single token → [1, EMBEDDING_DIM]", () => {
+    const embeddings = makeEmbeddings(10)
+    const input = T.fromArray(1, 1, [0])
+    const output = runEffect(embeddings.forward(input))
+    expectShape(output, [1, EMBEDDING_DIM])
+    expectFinite(output)
+  })
+
+  test("embed multiple tokens → [seqLen, EMBEDDING_DIM]", () => {
+    const embeddings = makeEmbeddings(10)
+    const input = T.fromArray(1, 3, [0, 1, 2])
+    const output = runEffect(embeddings.forward(input))
+    expectShape(output, [3, EMBEDDING_DIM])
+    expectFinite(output)
+  })
+
+  test("positional embeddings differ across positions", () => {
+    const embeddings = makeEmbeddings(10)
+    const input = T.fromArray(1, 3, [0, 0, 0])
+    const output = runEffect(embeddings.forward(input))
+    expectShape(output, [3, EMBEDDING_DIM])
+
+    const row0Val = T.get(output, 0, 0)
+    const row1Val = T.get(output, 1, 0)
+    const row2Val = T.get(output, 2, 0)
+
+    expect(row0Val).not.toBe(row1Val)
+    expect(row1Val).not.toBe(row2Val)
+    expect(row0Val).not.toBe(row2Val)
+  })
+
+  test("shape across sequence lengths 1..5", () => {
+    const embeddings = makeEmbeddings(10)
+    for (let seqLen = 1; seqLen <= 4; seqLen++) {
+      const tokens = Array.from({ length: seqLen }, (_, i) => i % 10)
+      const input = T.fromArray(1, seqLen, tokens)
+      const output = runEffect(embeddings.forward(input))
+      expectShape(output, [seqLen, EMBEDDING_DIM])
+      expectFinite(output)
+    }
+  })
+
+  test("max sequence length boundary", () => {
+    const embeddings = makeEmbeddings(10)
+    const tokens = Array.from({ length: MAX_SEQ_LEN }, (_, i) => i % 10)
+    const input = T.fromArray(1, MAX_SEQ_LEN, tokens)
+    const output = runEffect(embeddings.forward(input))
+    expectShape(output, [MAX_SEQ_LEN, EMBEDDING_DIM])
+    expectFinite(output)
+  })
+
+  test("backward updates token & positional embeddings", () => {
+    const embeddings = makeEmbeddings(10)
+    const input = T.fromArray(1, 3, [0, 1, 2])
+
+    const tokenBefore = T.clone(embeddings.tokenEmbeddings)
+    const positionalBefore = T.clone(embeddings.positionalEmbeddings)
+
+    runEffect(embeddings.forward(input))
+    const grad = T.ones(3, EMBEDDING_DIM)
+    runEffect(embeddings.backward(grad, 0.01))
+
+    expectNotClose(embeddings.tokenEmbeddings, tokenBefore)
+    expectNotClose(embeddings.positionalEmbeddings, positionalBefore)
+  })
+
+  test("parametersCount", () => {
+    const vocabSize = 10
+    const embeddings = makeEmbeddings(vocabSize)
+    const expectedCount = vocabSize * EMBEDDING_DIM + MAX_SEQ_LEN * EMBEDDING_DIM
+    expect(embeddings.parametersCount).toBe(expectedCount)
+  })
+})

package/tests/ts/errors.test.ts

@@ -0,0 +1,36 @@
+import { describe, expect, test } from "bun:test"
+import { TrainingError } from "../../src/errors"
+import { ShapeError } from "../../src/tensor/ops"
+import { DatasetLoadError, DatasetParseError } from "../../src/data/Dataset"
+
+describe("TrainingError.fromUnknown", () => {
+  test("maps ShapeError to TrainingShapeError", () => {
+    const err = TrainingError.fromUnknown(new ShapeError("bad shape"))
+    expect(err._tag).toBe("TrainingShapeError")
+    expect((err as any).cause.message).toBe("bad shape")
+  })
+
+  test("maps DatasetLoadError to TrainingDatasetError", () => {
+    const loadErr = new DatasetLoadError({ path: "x", error: new Error("io") })
+    const err = TrainingError.fromUnknown(loadErr)
+    expect(err._tag).toBe("TrainingDatasetError")
+    expect((err as any).cause).toBe(loadErr)
+  })
+
+  test("maps DatasetParseError to TrainingDatasetError", () => {
+    const parseErr = new DatasetParseError({ path: "x", error: new Error("parse") })
+    const err = TrainingError.fromUnknown(parseErr)
+    expect(err._tag).toBe("TrainingDatasetError")
+  })
+
+  test("passes through existing TrainingError", () => {
+    const existing = TrainingError.optimizer("boom")
+    const err = TrainingError.fromUnknown(existing)
+    expect(err).toBe(existing)
+  })
+
+  test("wraps unknown errors in TrainingUnknownError", () => {
+    const err = TrainingError.fromUnknown(new Error("boom"))
+    expect(err._tag).toBe("TrainingUnknownError")
+  })
+})

package/tests/ts/feed_forward.test.ts

@@ -0,0 +1,74 @@
+import { describe, test, expect } from "bun:test"
+import { runEffect } from "./support/runEffect"
+import { expectShape, expectNotClose, expectFinite } from "./support/tensorMatchers"
+import { makeFeedForward } from "./support/factories"
+import * as T from "../../src/tensor/Tensor2D"
+import { EMBEDDING_DIM, HIDDEN_DIM } from "../../src/config"
+
+describe("FeedForward", () => {
+  test("forward shape matches input", () => {
+    const ff = makeFeedForward()
+    const input = T.ones(3, EMBEDDING_DIM)
+    const output = runEffect(ff.forward(input))
+    expectShape(output, [3, EMBEDDING_DIM])
+  })
+
+  test("shape across sequence lengths 1..5", () => {
+    const ff = makeFeedForward()
+    for (let seqLen = 1; seqLen <= 4; seqLen++) {
+      const input = T.ones(seqLen, EMBEDDING_DIM)
+      const output = runEffect(ff.forward(input))
+      expectShape(output, [seqLen, EMBEDDING_DIM])
+    }
+  })
+
+  test("output contains finite values", () => {
+    const ff = makeFeedForward()
+    const input = T.ones(3, EMBEDDING_DIM)
+    const output = runEffect(ff.forward(input))
+    expectFinite(output)
+  })
+
+  test("backward returns gradient with correct shape", () => {
+    const ff = makeFeedForward()
+    const input = T.ones(3, EMBEDDING_DIM)
+    runEffect(ff.forward(input))
+    const grads = T.ones(3, EMBEDDING_DIM)
+    const gradInput = runEffect(ff.backward(grads, 0.01))
+    expectShape(gradInput, [3, EMBEDDING_DIM])
+  })
+
+  test("backward output differs from forward output", () => {
+    const ff = makeFeedForward()
+    const input = T.ones(3, EMBEDDING_DIM)
+    const output = runEffect(ff.forward(input))
+    const grads = T.ones(3, EMBEDDING_DIM)
+    const gradInput = runEffect(ff.backward(grads, 0.01))
+    expectNotClose(output, gradInput)
+  })
+
+  test("backward updates w1/b1/w2/b2", () => {
+    const ff = makeFeedForward()
+    const w1Before = T.clone(ff.w1)
+    const b1Before = T.clone(ff.b1)
+    const w2Before = T.clone(ff.w2)
+    const b2Before = T.clone(ff.b2)
+
+    const input = T.ones(3, EMBEDDING_DIM)
+    runEffect(ff.forward(input))
+    const grads = T.ones(3, EMBEDDING_DIM)
+    runEffect(ff.backward(grads, 0.01))
+
+    expectNotClose(ff.w1, w1Before)
+    expectNotClose(ff.b1, b1Before)
+    expectNotClose(ff.w2, w2Before)
+    expectNotClose(ff.b2, b2Before)
+  })
+
+  test("parametersCount", () => {
+    const ff = makeFeedForward()
+    const expected =
+      EMBEDDING_DIM * HIDDEN_DIM + HIDDEN_DIM + HIDDEN_DIM * EMBEDDING_DIM + EMBEDDING_DIM
+    expect(ff.parametersCount).toBe(expected)
+  })
+})

package/tests/ts/initNormal.test.ts

@@ -0,0 +1,41 @@
+import { describe, expect, it } from "bun:test"
+import { initNormal } from "../../src/tensor/ops"
+import { seeded } from "../../src/tensor/random"
+
+describe("initNormal with seeded RNG", () => {
+  it("produces deterministic tensors with same seed", () => {
+    const rng1 = seeded(42)
+    const rng2 = seeded(42)
+
+    const t1 = initNormal(4, 4, 0, 1, rng1)
+    const t2 = initNormal(4, 4, 0, 1, rng2)
+
+    expect(Array.from(t1.data)).toEqual(Array.from(t2.data))
+  })
+
+  it("produces different tensors with different seeds", () => {
+    const rng1 = seeded(42)
+    const rng2 = seeded(43)
+
+    const t1 = initNormal(4, 4, 0, 1, rng1)
+    const t2 = initNormal(4, 4, 0, 1, rng2)
+
+    expect(Array.from(t1.data)).not.toEqual(Array.from(t2.data))
+  })
+
+  it("respects mean and std parameters", () => {
+    const rng = seeded(1337)
+    const mean = 5
+    const std = 0.1
+    const t = initNormal(100, 100, mean, std, rng)
+
+    let sum = 0
+    for (let i = 0; i < t.data.length; i++) {
+      sum += t.data[i]
+    }
+    const actualMean = sum / t.data.length
+
+    // Mean should be close to target (within ~3 std errors)
+    expect(Math.abs(actualMean - mean)).toBeLessThan(0.1)
+  })
+})

package/tests/ts/layer_norm.test.ts

@@ -0,0 +1,96 @@
+import { describe, test, expect } from "bun:test"
+import { runEffect } from "./support/runEffect"
+import { expectShape, expectNotClose, expectFinite } from "./support/tensorMatchers"
+import { makeLayerNorm } from "./support/factories"
+import * as T from "../../src/tensor/Tensor2D"
+import { EMBEDDING_DIM } from "../../src/config"
+
+const rowMean = (t: T.Tensor2D, row: number): number => {
+  let sum = 0
+  for (let j = 0; j < t.cols; j++) sum += T.get(t, row, j)
+  return sum / t.cols
+}
+
+const rowVariance = (t: T.Tensor2D, row: number): number => {
+  const mean = rowMean(t, row)
+  let sumSq = 0
+  for (let j = 0; j < t.cols; j++) {
+    const diff = T.get(t, row, j) - mean
+    sumSq += diff * diff
+  }
+  return sumSq / t.cols
+}
+
+const makeVariedInput = (): T.Tensor2D => {
+  const input = T.zeros(3, EMBEDDING_DIM)
+  for (let i = 0; i < input.data.length; i++) {
+    input.data[i] = (i % 7) * 0.5 - 1.5
+  }
+  return input
+}
+
+describe("LayerNorm", () => {
+  test("forward shape preserved", () => {
+    const ln = makeLayerNorm()
+    const input = T.ones(3, EMBEDDING_DIM)
+    const output = runEffect(ln.forward(input))
+    expectShape(output, [3, EMBEDDING_DIM])
+  })
+
+  test("per-row mean approximately zero", () => {
+    const ln = makeLayerNorm()
+    const input = makeVariedInput()
+    const output = runEffect(ln.forward(input))
+
+    for (let row = 0; row < output.rows; row++) {
+      const mean = rowMean(output, row)
+      expect(Math.abs(mean)).toBeLessThan(1e-5)
+    }
+  })
+
+  test("per-row variance approximately one", () => {
+    const ln = makeLayerNorm()
+    const input = makeVariedInput()
+    const output = runEffect(ln.forward(input))
+
+    for (let row = 0; row < output.rows; row++) {
+      const variance = rowVariance(output, row)
+      expect(Math.abs(variance - 1)).toBeLessThan(1e-4)
+    }
+  })
+
+  test("output contains finite values", () => {
+    const ln = makeLayerNorm()
+    const input = makeVariedInput()
+    const output = runEffect(ln.forward(input))
+    expectFinite(output)
+  })
+
+  test("backward shape preserved", () => {
+    const ln = makeLayerNorm()
+    const input = makeVariedInput()
+    runEffect(ln.forward(input))
+    const dOut = T.ones(3, EMBEDDING_DIM)
+    const grad = runEffect(ln.backward(dOut, 0.01))
+    expectShape(grad, [3, EMBEDDING_DIM])
+  })
+
+  test("backward updates gamma and beta", () => {
+    const ln = makeLayerNorm()
+    const input = makeVariedInput()
+    const gammaBefore = T.clone(ln.gamma)
+    const betaBefore = T.clone(ln.beta)
+
+    runEffect(ln.forward(input))
+    const dOut = T.ones(3, EMBEDDING_DIM)
+    runEffect(ln.backward(dOut, 0.01))
+
+    expectNotClose(ln.gamma, gammaBefore)
+    expectNotClose(ln.beta, betaBefore)
+  })
+
+  test("parametersCount", () => {
+    const ln = makeLayerNorm()
+    expect(ln.parametersCount).toBe(2 * EMBEDDING_DIM)
+  })
+})