effect-gpt 0.1.0

package/src/services/SeedLayer.ts

@@ -0,0 +1,39 @@
+import * as Context from "effect/Context"
+import * as Layer from "effect/Layer"
+import * as Effect from "effect/Effect"
+import { seeded, systemRng, type Rng } from "../tensor/random"
+
+export interface SeedService {
+  /** Shared RNG stream used for deterministic initialization. */
+  readonly rng: Rng
+  /** Creates a new SeedService derived from the current stream (for isolation when desired). */
+  readonly fork: () => SeedService
+}
+
+export interface SeedServiceId {
+  readonly SeedService: unique symbol
+}
+
+export const Seed = Context.GenericTag<SeedServiceId, SeedService>("SeedService")
+
+const makeSeedService = (seed?: number): SeedService => {
+  const rng = seed === undefined ? systemRng() : seeded(seed)
+
+  return {
+    rng,
+    fork: () => {
+      // Advance the stream to produce a derived seed; fall back to system entropy when non-deterministic.
+      const nextSeed = Math.floor(rng.next() * 0xffffffff)
+      return makeSeedService(seed === undefined ? undefined : nextSeed)
+    }
+  }
+}
+
+/**
+ * Layer that provides a shared SeedService. Passing a seed yields deterministic initialization; omitting it uses
+ * nondeterministic Math.random.
+ */
+export const SeedLayer = (seed?: number): Layer.Layer<SeedServiceId> => Layer.succeed(Seed, makeSeedService(seed))
+
+/** Effect helper to grab the current RNG from context. */
+export const useSeedRng = (): Effect.Effect<Rng, never, SeedServiceId> => Effect.map(Seed, (service) => service.rng)

package/src/services/index.ts

@@ -0,0 +1,32 @@
+export type { LogLevel, LoggerService, LoggerServiceId } from "./Logger"
+export {
+  Logger,
+  ConsoleLoggerLive,
+  TerminalLoggerLive,
+  NullLoggerLive,
+  SilentLoggerLive,
+  PrettyLoggerLive,
+  log,
+  debug,
+  info,
+  warn,
+  error
+} from "./Logger"
+
+export type { RandomService, RandomServiceId } from "./Random"
+export { Random, SeededRandomLive, SystemRandomLive, next, nextGaussian, nextInt, fork } from "./Random"
+
+export type { SeedService, SeedServiceId } from "./SeedLayer"
+export { Seed, SeedLayer, useSeedRng } from "./SeedLayer"
+
+export type { Counter, Gauge, Histogram, TimingResult, MetricsService, MetricsSnapshot, MetricsServiceId } from "./Metrics"
+export {
+  Metrics,
+  InMemoryMetricsLive,
+  NoOpMetricsLive,
+  counter,
+  gauge,
+  histogram,
+  timed,
+  snapshot
+} from "./Metrics"

package/src/tensor/Tensor2D.ts

@@ -0,0 +1,42 @@
+export interface Tensor2D {
+  readonly rows: number
+  readonly cols: number
+  readonly data: Float32Array
+}
+
+export const make = (rows: number, cols: number, data: Float32Array): Tensor2D => {
+  if (data.length !== rows * cols) {
+    throw new Error(`Data length ${data.length} does not match shape ${rows}x${cols}`)
+  }
+  return { rows, cols, data }
+}
+
+export const zeros = (rows: number, cols: number): Tensor2D => {
+  return { rows, cols, data: new Float32Array(rows * cols) }
+}
+
+export const ones = (rows: number, cols: number): Tensor2D => {
+  const data = new Float32Array(rows * cols)
+  data.fill(1)
+  return { rows, cols, data }
+}
+
+export const clone = (t: Tensor2D): Tensor2D => {
+  return { rows: t.rows, cols: t.cols, data: new Float32Array(t.data) }
+}
+
+export const get = (t: Tensor2D, row: number, col: number): number => {
+  return t.data[row * t.cols + col]
+}
+
+export const set = (t: Tensor2D, row: number, col: number, value: number): void => {
+  t.data[row * t.cols + col] = value
+}
+
+export const fromArray = (rows: number, cols: number, arr: ArrayLike<number>): Tensor2D => {
+  const data = new Float32Array(arr)
+  if (data.length !== rows * cols) {
+    throw new Error(`Array length ${data.length} does not match shape ${rows}x${cols}`)
+  }
+  return { rows, cols, data }
+}

package/src/tensor/ops.ts

@@ -0,0 +1,371 @@
+import * as Effect from "effect/Effect"
+import type { Tensor2D } from "./Tensor2D"
+import * as T from "./Tensor2D"
+import type { Rng } from "./random"
+import type { RandomServiceId } from "../services/Random"
+import { Random } from "../services/Random"
+
+export class ShapeError extends Error {
+  readonly _tag = "ShapeError"
+  constructor(message: string) {
+    super(message)
+    this.name = "ShapeError"
+  }
+}
+
+export const matMul = (a: Tensor2D, b: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (a.cols !== b.rows) {
+      throw new ShapeError(`matMul: a.cols (${a.cols}) !== b.rows (${b.rows})`)
+    }
+    const result = T.zeros(a.rows, b.cols)
+    for (let i = 0; i < a.rows; i++) {
+      for (let j = 0; j < b.cols; j++) {
+        let sum = 0
+        for (let k = 0; k < a.cols; k++) {
+          sum += T.get(a, i, k) * T.get(b, k, j)
+        }
+        T.set(result, i, j, sum)
+      }
+    }
+    return result
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const add = (a: Tensor2D, b: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (a.rows !== b.rows || a.cols !== b.cols) {
+      throw new ShapeError(`add: shapes (${a.rows},${a.cols}) and (${b.rows},${b.cols}) do not match`)
+    }
+    const data = new Float32Array(a.data.length)
+    for (let i = 0; i < data.length; i++) {
+      data[i] = a.data[i] + b.data[i]
+    }
+    return T.make(a.rows, a.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const sub = (a: Tensor2D, b: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (a.rows !== b.rows || a.cols !== b.cols) {
+      throw new ShapeError(`sub: shapes (${a.rows},${a.cols}) and (${b.rows},${b.cols}) do not match`)
+    }
+    const data = new Float32Array(a.data.length)
+    for (let i = 0; i < data.length; i++) {
+      data[i] = a.data[i] - b.data[i]
+    }
+    return T.make(a.rows, a.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const mul = (a: Tensor2D, b: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (a.rows !== b.rows || a.cols !== b.cols) {
+      throw new ShapeError(`mul: shapes (${a.rows},${a.cols}) and (${b.rows},${b.cols}) do not match`)
+    }
+    const data = new Float32Array(a.data.length)
+    for (let i = 0; i < data.length; i++) {
+      data[i] = a.data[i] * b.data[i]
+    }
+    return T.make(a.rows, a.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const div = (a: Tensor2D, b: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (a.rows !== b.rows || a.cols !== b.cols) {
+      throw new ShapeError(`div: shapes (${a.rows},${a.cols}) and (${b.rows},${b.cols}) do not match`)
+    }
+    const data = new Float32Array(a.data.length)
+    for (let i = 0; i < data.length; i++) {
+      data[i] = a.data[i] / b.data[i]
+    }
+    return T.make(a.rows, a.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const addScalar = (t: Tensor2D, scalar: number): Tensor2D => {
+  const data = new Float32Array(t.data.length)
+  for (let i = 0; i < data.length; i++) {
+    data[i] = t.data[i] + scalar
+  }
+  return T.make(t.rows, t.cols, data)
+}
+
+export const mulScalar = (t: Tensor2D, scalar: number): Tensor2D => {
+  const data = new Float32Array(t.data.length)
+  for (let i = 0; i < data.length; i++) {
+    data[i] = t.data[i] * scalar
+  }
+  return T.make(t.rows, t.cols, data)
+}
+
+export const addRowBias = (matrix: Tensor2D, bias: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (bias.rows !== 1 || bias.cols !== matrix.cols) {
+      throw new ShapeError(`addRowBias: bias shape (${bias.rows},${bias.cols}) incompatible with matrix cols ${matrix.cols}`)
+    }
+    const data = new Float32Array(matrix.data.length)
+    for (let i = 0; i < matrix.rows; i++) {
+      for (let j = 0; j < matrix.cols; j++) {
+        data[i * matrix.cols + j] = T.get(matrix, i, j) + bias.data[j]
+      }
+    }
+    return T.make(matrix.rows, matrix.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const meanRows = (t: Tensor2D): Tensor2D => {
+  const data = new Float32Array(t.rows)
+  for (let i = 0; i < t.rows; i++) {
+    let sum = 0
+    for (let j = 0; j < t.cols; j++) {
+      sum += T.get(t, i, j)
+    }
+    data[i] = sum / t.cols
+  }
+  return T.make(t.rows, 1, data)
+}
+
+export const sumCols = (t: Tensor2D): Tensor2D => {
+  const data = new Float32Array(t.cols)
+  for (let j = 0; j < t.cols; j++) {
+    let sum = 0
+    for (let i = 0; i < t.rows; i++) {
+      sum += T.get(t, i, j)
+    }
+    data[j] = sum
+  }
+  return T.make(1, t.cols, data)
+}
+
+export const meanCols = (t: Tensor2D): Tensor2D => {
+  const data = new Float32Array(t.cols)
+  const scale = 1 / t.rows
+  for (let j = 0; j < t.cols; j++) {
+    let sum = 0
+    for (let i = 0; i < t.rows; i++) {
+      sum += T.get(t, i, j)
+    }
+    data[j] = sum * scale
+  }
+  return T.make(1, t.cols, data)
+}
+
+export const stdRows = (t: Tensor2D): Tensor2D => {
+  const means = meanRows(t)
+  const data = new Float32Array(t.rows)
+  for (let i = 0; i < t.rows; i++) {
+    const mean = means.data[i]
+    let sumSq = 0
+    for (let j = 0; j < t.cols; j++) {
+      const diff = T.get(t, i, j) - mean
+      sumSq += diff * diff
+    }
+    data[i] = Math.sqrt(sumSq / t.cols)
+  }
+  return T.make(t.rows, 1, data)
+}
+
+export const varRows = (t: Tensor2D): Tensor2D => {
+  const means = meanRows(t)
+  const data = new Float32Array(t.rows)
+  for (let i = 0; i < t.rows; i++) {
+    const mean = means.data[i]
+    let sumSq = 0
+    for (let j = 0; j < t.cols; j++) {
+      const diff = T.get(t, i, j) - mean
+      sumSq += diff * diff
+    }
+    data[i] = sumSq / t.cols
+  }
+  return T.make(t.rows, 1, data)
+}
+
+export const mapScalar = (t: Tensor2D, fn: (val: number) => number): Tensor2D => {
+  const data = new Float32Array(t.data.length)
+  for (let i = 0; i < data.length; i++) {
+    data[i] = fn(t.data[i])
+  }
+  return T.make(t.rows, t.cols, data)
+}
+
+export const softmaxRows = (t: Tensor2D): Tensor2D => {
+  const data = new Float32Array(t.data.length)
+  for (let i = 0; i < t.rows; i++) {
+    let maxVal = -Infinity
+    for (let j = 0; j < t.cols; j++) {
+      const val = T.get(t, i, j)
+      if (val > maxVal) maxVal = val
+    }
+    let sumExp = 0
+    for (let j = 0; j < t.cols; j++) {
+      const exp = Math.exp(T.get(t, i, j) - maxVal)
+      data[i * t.cols + j] = exp
+      sumExp += exp
+    }
+    for (let j = 0; j < t.cols; j++) {
+      data[i * t.cols + j] /= sumExp
+    }
+  }
+  return T.make(t.rows, t.cols, data)
+}
+
+export const transpose = (t: Tensor2D): Tensor2D => {
+  const data = new Float32Array(t.rows * t.cols)
+  for (let i = 0; i < t.rows; i++) {
+    for (let j = 0; j < t.cols; j++) {
+      data[j * t.rows + i] = T.get(t, i, j)
+    }
+  }
+  return T.make(t.cols, t.rows, data)
+}
+
+export const gatherRows = (embeddings: Tensor2D, tokenIds: ReadonlyArray<number>): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    const data = new Float32Array(tokenIds.length * embeddings.cols)
+    for (let i = 0; i < tokenIds.length; i++) {
+      const tokenId = tokenIds[i]
+      if (tokenId < 0 || tokenId >= embeddings.rows) {
+        throw new ShapeError(`gatherRows: tokenId ${tokenId} out of bounds [0, ${embeddings.rows})`)
+      }
+      for (let j = 0; j < embeddings.cols; j++) {
+        data[i * embeddings.cols + j] = T.get(embeddings, tokenId, j)
+      }
+    }
+    return T.make(tokenIds.length, embeddings.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const sliceRows = (t: Tensor2D, start: number, end: number): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (start < 0 || end > t.rows || start >= end) {
+      throw new ShapeError(`sliceRows: invalid range [${start}, ${end}) for tensor with ${t.rows} rows`)
+    }
+    const numRows = end - start
+    const data = new Float32Array(numRows * t.cols)
+    for (let i = 0; i < numRows; i++) {
+      for (let j = 0; j < t.cols; j++) {
+        data[i * t.cols + j] = T.get(t, start + i, j)
+      }
+    }
+    return T.make(numRows, t.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const relu = (t: Tensor2D): Tensor2D => {
+  const data = new Float32Array(t.data.length)
+  for (let i = 0; i < data.length; i++) {
+    data[i] = Math.max(0, t.data[i])
+  }
+  return T.make(t.rows, t.cols, data)
+}
+
+export const argmaxRows = (t: Tensor2D): ReadonlyArray<number> => {
+  const result: Array<number> = []
+  for (let i = 0; i < t.rows; i++) {
+    let maxIdx = 0
+    let maxVal = T.get(t, i, 0)
+    for (let j = 1; j < t.cols; j++) {
+      const val = T.get(t, i, j)
+      if (val > maxVal) {
+        maxVal = val
+        maxIdx = j
+      }
+    }
+    result.push(maxIdx)
+  }
+  return result
+}
+
+export const broadcastSubCol = (t: Tensor2D, col: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (col.cols !== 1 || col.rows !== t.rows) {
+      throw new ShapeError(`broadcastSubCol: col shape (${col.rows},${col.cols}) incompatible with tensor rows ${t.rows}`)
+    }
+    const data = new Float32Array(t.data.length)
+    for (let i = 0; i < t.rows; i++) {
+      for (let j = 0; j < t.cols; j++) {
+        data[i * t.cols + j] = T.get(t, i, j) - col.data[i]
+      }
+    }
+    return T.make(t.rows, t.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const broadcastDivCol = (t: Tensor2D, col: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (col.cols !== 1 || col.rows !== t.rows) {
+      throw new ShapeError(`broadcastDivCol: col shape (${col.rows},${col.cols}) incompatible with tensor rows ${t.rows}`)
+    }
+    const data = new Float32Array(t.data.length)
+    for (let i = 0; i < t.rows; i++) {
+      for (let j = 0; j < t.cols; j++) {
+        data[i * t.cols + j] = T.get(t, i, j) / col.data[i]
+      }
+    }
+    return T.make(t.rows, t.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const broadcastMulCol = (t: Tensor2D, col: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (col.cols !== 1 || col.rows !== t.rows) {
+      throw new ShapeError(`broadcastMulCol: col shape (${col.rows},${col.cols}) incompatible with tensor rows ${t.rows}`)
+    }
+    const data = new Float32Array(t.data.length)
+    for (let i = 0; i < t.rows; i++) {
+      for (let j = 0; j < t.cols; j++) {
+        data[i * t.cols + j] = T.get(t, i, j) * col.data[i]
+      }
+    }
+    return T.make(t.rows, t.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const broadcastMulRow = (t: Tensor2D, row: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (row.rows !== 1 || row.cols !== t.cols) {
+      throw new ShapeError(`broadcastMulRow: row shape (${row.rows},${row.cols}) incompatible with tensor cols ${t.cols}`)
+    }
+    const data = new Float32Array(t.data.length)
+    for (let i = 0; i < t.rows; i++) {
+      for (let j = 0; j < t.cols; j++) {
+        data[i * t.cols + j] = T.get(t, i, j) * row.data[j]
+      }
+    }
+    return T.make(t.rows, t.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const broadcastAddRow = (t: Tensor2D, row: Tensor2D): Effect.Effect<Tensor2D, ShapeError> =>
+  Effect.sync(() => {
+    if (row.rows !== 1 || row.cols !== t.cols) {
+      throw new ShapeError(`broadcastAddRow: row shape (${row.rows},${row.cols}) incompatible with tensor cols ${t.cols}`)
+    }
+    const data = new Float32Array(t.data.length)
+    for (let i = 0; i < t.rows; i++) {
+      for (let j = 0; j < t.cols; j++) {
+        data[i * t.cols + j] = T.get(t, i, j) + row.data[j]
+      }
+    }
+    return T.make(t.rows, t.cols, data)
+  }).pipe(Effect.catchAllDefect((e) => Effect.fail(e as ShapeError)))
+
+export const initNormal = (rows: number, cols: number, mean: number, std: number, rng: Rng): Tensor2D => {
+  const rand = () => rng.next()
+  const data = new Float32Array(rows * cols)
+  for (let i = 0; i < data.length; i++) {
+    let u1 = rand()
+    let u2 = rand()
+    while (u1 === 0) u1 = rand()
+    const z = Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2)
+    data[i] = mean + std * z
+  }
+  return T.make(rows, cols, data)
+}
+
+export const initNormalEffect = (
+  rows: number,
+  cols: number,
+  mean: number,
+  std: number
+): Effect.Effect<Tensor2D, never, RandomServiceId> =>
+  Effect.flatMap(Random, (random) =>
+    Effect.gen(function* () {
+      const data = new Float32Array(rows * cols)
+      for (let i = 0; i < data.length; i++) {
+        const value = yield* random.nextGaussian(mean, std)
+        data[i] = value
+      }
+      return T.make(rows, cols, data)
+    })
+  )

package/src/tensor/random.ts

@@ -0,0 +1,32 @@
+/**
+ * Deterministic random number generator for reproducible weight initialization.
+ * Uses mulberry32 algorithm for fast, seedable pseudo-random numbers.
+ */
+
+export interface Rng {
+  /** Returns a random number in [0, 1) */
+  next(): number
+}
+
+/**
+ * Creates a seeded RNG using the mulberry32 algorithm.
+ * Produces deterministic sequences for reproducible tests.
+ */
+export const seeded = (seed: number): Rng => {
+  let state = seed >>> 0
+
+  return {
+    next(): number {
+      state = (state + 0x6d2b79f5) >>> 0
+      let t = state
+      t = Math.imul(t ^ (t >>> 15), t | 1)
+      t ^= t + Math.imul(t ^ (t >>> 7), t | 61)
+      return ((t ^ (t >>> 14)) >>> 0) / 4294967296
+    }
+  }
+}
+
+/** Non-deterministic RNG wrapper around Math.random for convenience. */
+export const systemRng = (): Rng => ({
+  next: () => Math.random()
+})

package/src/tokenize/split.ts

@@ -0,0 +1,27 @@
+const ASCII_PUNCTUATION = /[!"#$%&'()*+,\-./:;<=>?@[\\\]^_`{|}~]/
+
+const isAsciiPunctuation = (ch: string): boolean =>
+  ch.length === 1 && ASCII_PUNCTUATION.test(ch)
+
+export const splitWordToTokens = (word: string): ReadonlyArray<string> => {
+  const tokens: Array<string> = []
+  let current = ""
+
+  for (const ch of word) {
+    if (isAsciiPunctuation(ch)) {
+      if (current.length > 0) {
+        tokens.push(current)
+        current = ""
+      }
+      tokens.push(ch)
+    } else {
+      current += ch
+    }
+  }
+
+  if (current.length > 0) {
+    tokens.push(current)
+  }
+
+  return tokens
+}

package/src/tokenize/tokenize.ts

@@ -0,0 +1,28 @@
+import * as Option from "effect/Option"
+import { splitWordToTokens } from "./split"
+import { Vocab } from "../vocab/Vocab"
+
+export const tokenize = (text: string, vocab: Vocab): ReadonlyArray<number> => {
+  const tokens: Array<number> = []
+  const words = text.split(/\s+/).filter((w) => w.length > 0)
+
+  for (const word of words) {
+    if (word === "</s>") {
+      const tokenId = vocab.encode(word)
+      if (Option.isSome(tokenId)) {
+        tokens.push(tokenId.value as number)
+      }
+      continue
+    }
+
+    const split = splitWordToTokens(word)
+    for (const part of split) {
+      const tokenId = vocab.encode(part)
+      if (Option.isSome(tokenId)) {
+        tokens.push(tokenId.value as number)
+      }
+    }
+  }
+
+  return tokens
+}

package/src/training/Adam.ts

@@ -0,0 +1,61 @@
+import type { Tensor2D } from "../tensor/Tensor2D"
+import * as T from "../tensor/Tensor2D"
+import { ShapeError } from "../tensor/ops"
+
+export class Adam {
+  readonly beta1 = 0.9
+  readonly beta2 = 0.999
+  readonly epsilon = 1e-8
+  timestep = 0
+  m: Tensor2D
+  v: Tensor2D
+
+  private constructor(rows: number, cols: number) {
+    this.m = T.zeros(rows, cols)
+    this.v = T.zeros(rows, cols)
+  }
+
+  static make(rows: number, cols: number): Adam {
+    return new Adam(rows, cols)
+  }
+
+  step(params: Tensor2D, grads: Tensor2D, lr: number): void {
+    if (params.rows !== grads.rows || params.cols !== grads.cols) {
+      throw new ShapeError(
+        `Adam.step: params shape (${params.rows},${params.cols}) != grads shape (${grads.rows},${grads.cols})`
+      )
+    }
+    if (this.m.rows !== params.rows || this.m.cols !== params.cols) {
+      throw new ShapeError(
+        `Adam.step: optimizer shape (${this.m.rows},${this.m.cols}) != params shape (${params.rows},${params.cols})`
+      )
+    }
+
+    this.timestep += 1
+    const beta1 = this.beta1
+    const beta2 = this.beta2
+    const oneMinusB1 = 1 - beta1
+    const oneMinusB2 = 1 - beta2
+
+    const mData = this.m.data
+    const vData = this.v.data
+    const pData = params.data
+    const gData = grads.data
+
+    for (let i = 0; i < gData.length; i++) {
+      const g = gData[i]
+      mData[i] = mData[i] * beta1 + g * oneMinusB1
+      vData[i] = vData[i] * beta2 + g * g * oneMinusB2
+    }
+
+    const mHatScale = 1 - Math.pow(beta1, this.timestep)
+    const vHatScale = 1 - Math.pow(beta2, this.timestep)
+
+    for (let i = 0; i < pData.length; i++) {
+      const mHat = mData[i] / mHatScale
+      const vHat = vData[i] / vHatScale
+      const update = mHat / (Math.sqrt(vHat) + this.epsilon)
+      pData[i] -= lr * update
+    }
+  }
+}

package/src/training/clip.ts

@@ -0,0 +1,16 @@
+import type { Tensor2D } from "../tensor/Tensor2D"
+
+export const clipGlobalL2 = (grads: Tensor2D, maxNorm: number): void => {
+  let sumSq = 0
+  for (let i = 0; i < grads.data.length; i++) {
+    const v = grads.data[i]
+    sumSq += v * v
+  }
+  const norm = Math.sqrt(sumSq)
+  if (norm > maxNorm) {
+    const scale = maxNorm / norm
+    for (let i = 0; i < grads.data.length; i++) {
+      grads.data[i] *= scale
+    }
+  }
+}

package/src/training/loss.ts

@@ -0,0 +1,35 @@
+import type { Tensor2D } from "../tensor/Tensor2D"
+import * as T from "../tensor/Tensor2D"
+import * as Ops from "../tensor/ops"
+
+export const softmaxRows = (logits: Tensor2D): Tensor2D => Ops.softmaxRows(logits)
+
+export const crossEntropyLoss = (probs: Tensor2D, targetIds: ReadonlyArray<number>): number => {
+  if (probs.rows !== targetIds.length) {
+    throw new Ops.ShapeError(`crossEntropyLoss: probs.rows (${probs.rows}) !== targetIds.length (${targetIds.length})`)
+  }
+  let loss = 0
+  for (let i = 0; i < probs.rows; i++) {
+    const idx = targetIds[i]
+    const prob = probs.data[i * probs.cols + idx]
+    const clamped = prob < 1e-15 ? 1e-15 : prob
+    loss -= Math.log(clamped)
+  }
+  return loss / targetIds.length
+}
+
+export const dLogits = (probs: Tensor2D, targetIds: ReadonlyArray<number>): Tensor2D => {
+  if (probs.rows !== targetIds.length) {
+    throw new Ops.ShapeError(`dLogits: probs.rows (${probs.rows}) !== targetIds.length (${targetIds.length})`)
+  }
+  const data = new Float32Array(probs.data)
+  for (let i = 0; i < probs.rows; i++) {
+    const idx = targetIds[i]
+    data[i * probs.cols + idx] -= 1
+  }
+  const scale = 1 / targetIds.length
+  for (let i = 0; i < data.length; i++) {
+    data[i] *= scale
+  }
+  return T.make(probs.rows, probs.cols, data)
+}