Initial commit

.gitignore

@@ -0,0 +1,4 @@
+node_modules
+build
+data
+models

eslint.config.mjs

@@ -0,0 +1,59 @@
+import { FlatCompat } from '@eslint/eslintrc'
+import js from '@eslint/js'
+import stylistic from '@stylistic/eslint-plugin'
+import stylisticts from '@stylistic/eslint-plugin-ts'
+import ts from '@typescript-eslint/eslint-plugin'
+import parser from '@typescript-eslint/parser'
+import path from 'node:path'
+import { fileURLToPath } from 'node:url'
+
+export default [
+  ...new FlatCompat({
+    baseDirectory: path.dirname(fileURLToPath(import.meta.url)),
+    recommendedConfig: js.configs.recommended,
+    allConfig: js.configs.all,
+  })
+    .extends(
+      'eslint:recommended',
+      'plugin:@typescript-eslint/eslint-recommended',
+      'plugin:@typescript-eslint/recommended',
+    ),
+  stylistic.configs['recommended-flat'],
+  {
+    files: ['src/**/*.ts', 'test/**/*.ts'],
+    languageOptions: { parser },
+    plugins: {
+      '@typescript-eslint': ts,
+      '@stylistic': stylistic,
+      '@stylistic/ts': stylisticts,
+    },
+    rules: {
+      'arrow-body-style': 'error',
+      'no-extra-parens': ['error', 'all'],
+      'no-unexpected-multiline': 'off',
+      'sort-imports': ['error', { ignoreDeclarationSort: true }],
+      '@stylistic/brace-style': ['error', '1tbs', { allowSingleLine: true }],
+      '@stylistic/quotes': ['error', 'single'],
+      '@stylistic/semi': ['error', 'always'],
+      '@stylistic/ts/semi': ['error', 'always'],
+      '@stylistic/indent': 'error',
+      '@stylistic/arrow-parens': ['error', 'as-needed', { requireForBlockBody: false }],
+      '@stylistic/max-statements-per-line': 'off',
+      '@stylistic/member-delimiter-style': 'off',
+      '@stylistic/ts/member-delimiter-style': 'error',
+      '@stylistic/no-mixed-operators': 'off',
+      '@typescript-eslint/no-unused-vars': [
+        'error',
+        {
+          args: 'all',
+          argsIgnorePattern: '^_',
+          caughtErrors: 'all',
+          caughtErrorsIgnorePattern: '^_',
+          destructuredArrayIgnorePattern: '^_',
+          varsIgnorePattern: '^_',
+          ignoreRestSiblings: true,
+        },
+      ],
+      '@typescript-eslint/no-explicit-any': 'off',
+    },
+  }]

jest.config.ts

@@ -0,0 +1,8 @@
+import type { Config } from 'jest'
+
+export default {
+  testEnvironment: 'node',
+  transform: {
+    '^.+.tsx?$': ['ts-jest', { tsconfig: 'test/tsconfig.json' }],
+  },
+} as Config

package.json

@@ -0,0 +1,40 @@
+{
+  "name": "nn",
+  "version": "1.0.0",
+  "type": "module",
+  "main": "index.js",
+  "author": "Freywar Ulvnaudgari <freywar.ulvnaudgari@gmail.com>",
+  "license": "MIT",
+  "scripts": {
+    "build": "rm -rf build && tsc",
+    "start": "yarn build && systemd-inhibit --what=idle:sleep --why='Training a model' node --import=extensionless/register build/index.js",
+    "start:debug": "yarn build && node --inspect-brk --import=extensionless/register build/index.js",
+    "start:profile": "yarn build && node --prof --import=extensionless/register build/index.js",
+    "test": "jest",
+    "test:debug": "node --inspect-brk node_modules/.bin/jest --runInBand",
+    "lint": "eslint",
+    "lint:fix": "eslint --fix"
+  },
+  "devDependencies": {
+    "@eslint/eslintrc": "^3.1.0",
+    "@eslint/js": "^9.8.0",
+    "@jest/globals": "^29.7.0",
+    "@stylistic/eslint-plugin": "^2.6.1",
+    "@stylistic/eslint-plugin-ts": "^2.6.4",
+    "@typescript-eslint/eslint-plugin": "^8.0.0",
+    "@typescript-eslint/parser": "^8.0.0",
+    "eslint": "^9.8.0",
+    "extensionless": "^1.9.9",
+    "jest": "^29.7.0",
+    "ts-jest": "^29.2.3",
+    "ts-node": "^10.9.2",
+    "typescript": "^5.5.4"
+  },
+  "dependencies": {
+    "@webgpu/glslang": "^0.0.15",
+    "@webgpu/types": "^0.1.60",
+    "ts-command-line-args": "^2.5.1",
+    "webgpu": "^0.2.11",
+    "yargs": "^17.7.2"
+  }
+}

shader/argmax.wgsl

@@ -0,0 +1,37 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+struct FVectorBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+struct UVectorBatch {
+  size: vec4<u32>,
+  data: array<u32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> ps: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: UVectorBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size.x = ps.size.y;
+  result.size.z = ps.size.z;
+
+  var mi = 0u;
+  var mp = 0f;
+  var pi = global_id.z * ps.size.y * ps.size.x + global_id.x * ps.size.x;
+  for (var i = 0u; i < ps.size.x; i++) {
+    if (ps.data[pi] > mp) {
+      mi = i;
+      mp = ps.data[pi];
+    }
+    pi++;
+  }
+  result.data[global_id.z * result.size.x + global_id.x] = mi;
+}

shader/attw.wgsl

@@ -0,0 +1,33 @@
+struct PackedMatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> e: PackedMatrixBatch;
+@group(0) @binding(1)
+var<storage, read> tw: PackedMatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> aw: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  aw.size.x = tw.size.y;
+  aw.size.y = e.size.y;
+  aw.size.z = e.size.z;
+
+  var s = 0f;
+  let vc = e.size.x / 4;
+  let lo = (global_id.z * e.size.y + global_id.y) * vc;
+  let ro = (global_id.z * tw.size.y + global_id.x) * vc;
+  for (var i = 0u; i < vc; i++) {
+    s += dot(e.data[lo + i], tw.data[ro + i]);
+  }
+
+  aw.data[(global_id.z * aw.size.y + global_id.y) * aw.size.x + global_id.x] = s;
+}

shader/avg.wgsl

@@ -0,0 +1,32 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+struct Matrix {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: Matrix;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = m.size;
+  result.size.z = 1;
+
+  var vc = result.size.x / 4;
+  var resi = global_id.y * vc + global_id.x;
+  result.data[resi] = vec4(0f);
+
+  var mi = resi;
+  for (var z = 0u; z < m.size.z; z++) {
+    result.data[resi] += m.data[mi] / vec4(f32(m.size.z));
+    mi += m.size.y * vc;
+  }
+}

shader/base.wgsl

@@ -0,0 +1,18 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+@group(0) @binding(3)
+var<storage, read_write> m: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  for (var z = 0u; z < m.size.z; z++) {
+    for (var y = 0u; y < m.size.y; y++) {
+      for (var x = 0u; x < m.size.x; x++) {
+        m.data[z * m.size.y * m.size.x + y * m.size.x + x] = 0f;
+      }
+    }
+  }
+}

shader/bscale.wgsl

@@ -0,0 +1,26 @@
+alias number = f32;
+
+struct ScalarBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: ScalarBatch;
+@group(0) @binding(1)
+var<storage, read> r_or_result: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = r_or_result.size;
+
+  var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x;
+  result.data[resi] = l.data[global_id.z] * r_or_result.data[resi];
+}

shader/cdiv.wgsl

@@ -0,0 +1,21 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> r_or_result: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = l.size;
+
+  var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x;
+  result.data[resi] = l.data[resi] / r_or_result.data[resi];
+}

shader/cmul.wgsl

@@ -0,0 +1,21 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> r_or_result: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = l.size;
+
+  var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x;
+  result.data[resi] = l.data[resi] * r_or_result.data[resi];
+}

shader/cscale.wgsl

@@ -0,0 +1,26 @@
+alias number = f32;
+
+struct Scalar {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: Scalar;
+@group(0) @binding(1)
+var<storage, read> r_or_result: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = r_or_result.size;
+
+  var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x;
+  result.data[resi] = l.data[0] * r_or_result.data[resi];
+}

shader/csub.wgsl

@@ -0,0 +1,21 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> r_or_result: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = l.size;
+
+  var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x;
+  result.data[resi] = l.data[resi] - r_or_result.data[resi];
+}

shader/csum.wgsl

@@ -0,0 +1,21 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> r_or_result: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = l.size;
+
+  var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x;
+  result.data[resi] = l.data[resi] + r_or_result.data[resi];
+}

shader/dldo.wgsl

@@ -0,0 +1,34 @@
+alias number = f32;
+
+struct Scalar {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+struct UVectorBatch {
+  size: vec4<u32>,
+  data: array<u32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> v: UVectorBatch;
+@group(0) @binding(2)
+var<storage, read> lr: Scalar;
+@group(0) @binding(3)
+var<storage, read_write> grs: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  grs.size = m.size;
+
+  var vi = global_id.z * v.size.x + global_id.y;
+  var gi = vi * grs.size.x + global_id.x;
+  grs.data[gi] = lr.data[0] * (m.data[gi] - select(0f, 1f, v.data[vi] == global_id.x));
+}

shader/dlds.wgsl

@@ -0,0 +1,48 @@
+struct Scalar {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> ds: Scalar;
+@group(0) @binding(1)
+var<storage, read> wa: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read> v: MatrixBatch;
+@group(0) @binding(3)
+var<storage, read> otdldo: MatrixBatch;
+@group(0) @binding(4)
+var<storage, read_write> dlds: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  dlds.size = wa.size;
+
+  var s = 0f;
+  var li = (global_id.z * v.size.y + global_id.x) * v.size.x;
+  var ri = global_id.z * otdldo.size.x * otdldo.size.y + global_id.y;
+
+  for (var i = 0u; i < v.size.x; i++) {
+    var cs = 0f;
+    var lj = global_id.z * wa.size.y * wa.size.x + global_id.y * wa.size.x;
+    var mj = global_id.z * v.size.x * v.size.y + i;
+    for (var j = 0u; j < wa.size.x; j++) {
+      cs += wa.data[lj] * v.data[mj];
+      lj++;
+      mj += v.size.x;
+    }
+
+    s += (v.data[li] - cs) * otdldo.data[ri];
+
+    li++;
+    ri += otdldo.size.x;
+  }
+
+  let resi = (global_id.z * dlds.size.y + global_id.y) * dlds.size.x + global_id.x;
+  dlds.data[resi] = ds.data[0] * wa.data[resi] * s;
+}

shader/dots.wgsl

@@ -0,0 +1,43 @@
+struct Scalar {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+struct Matrix {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> ebs: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> ows: Matrix;
+@group(0) @binding(2)
+var<storage, read> tmp: Scalar;
+@group(0) @binding(3)
+var<storage, read_write> dts: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  dts.size.x = ows.size.y;
+  dts.size.y = ebs.size.y;
+  dts.size.z = ebs.size.z;
+
+  let vc = ebs.size.x / 4;
+  let eo = (global_id.z * ebs.size.y + global_id.y) * vc;
+  let oo = 4 * global_id.x * vc;
+  var s = vec4(0f);
+  for (var i = 0u; i < vc; i++) {
+    let eb = ebs.data[eo + i];
+    s.x += dot(eb, ows.data[oo + 0 * vc + i]);
+    s.y += dot(eb, ows.data[oo + 1 * vc + i]);
+    s.z += dot(eb, ows.data[oo + 2 * vc + i]);
+    s.w += dot(eb, ows.data[oo + 3 * vc + i]);
+  }
+  dts.data[(global_id.z * dts.size.y + global_id.y) * dts.size.x / 4 + global_id.x] = tmp.data[0] * s;
+}

shader/embed.wgsl

@@ -0,0 +1,30 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> ilogits: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> inp: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read> pos: MatrixBatch;
+@group(0) @binding(3)
+var<storage, read_write> r: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  r.size = pos.size;
+
+  var s = 0f;
+  var li = (global_id.z * ilogits.size.y + global_id.y) * ilogits.size.x;
+  var ri = global_id.z * inp.size.x * inp.size.y + global_id.x;
+  for (var i = 0u; i < ilogits.size.x; i++) {
+    s += ilogits.data[li] * inp.data[ri];
+    li++;
+    ri += inp.size.x;
+  }
+
+  var resi = (global_id.z * r.size.y + global_id.y) * r.size.x + global_id.x;
+  r.data[resi] = s + pos.data[resi];
+}

shader/gradinp.wgsl

@@ -0,0 +1,47 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> ilogits: MatrixBatch;
+@group(0) @binding(0)
+var<storage, read> tilogits: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> dldsq: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read> dldsk: MatrixBatch;
+@group(0) @binding(3)
+var<storage, read> totdldowa: MatrixBatch;
+@group(0) @binding(4)
+var<storage, read> tqry: MatrixBatch;
+@group(0) @binding(5)
+var<storage, read> tkey: MatrixBatch;
+@group(0) @binding(6)
+var<storage, read> tval: MatrixBatch;
+@group(0) @binding(7)
+var<storage, read_write> ginp: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  ginp.size.x = tval.size.x;
+  ginp.size.y = ilogits.size.x;
+  ginp.size.z = ilogits.size.z;
+
+  var ts = 0f;
+  var li = global_id.z * ilogits.size.y * ilogits.size.x + global_id.y;
+  for (var i = 0u; i < ilogits.size.y; i++) {
+    var s = 0f;
+    var lj = global_id.z * totdldowa.size.x * totdldowa.size.y + i * totdldowa.size.x;
+    var rj = global_id.z * tval.size.x * tval.size.y + global_id.x;
+    for (var j = 0u; j < totdldowa.size.x; j++) {
+      s += totdldowa.data[lj] * tval.data[rj] + dldsq.data[lj] * tkey.data[rj] + dldsk.data[lj] * tqry.data[rj];
+      lj++;
+      rj += tval.size.x;
+    }
+    ts += ilogits.data[li] * s;
+    li += ilogits.size.x;
+  }
+
+  ginp.data[global_id.z * ginp.size.y * ginp.size.x + global_id.y * ginp.size.x + global_id.x] = ts;
+}

shader/gradtattw.wgsl

@@ -0,0 +1,28 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> e: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> dlds: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> gtattw: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  gtattw.size.z = e.size.z;
+  gtattw.size.y = e.size.x;
+  gtattw.size.x = e.size.x;
+
+  var s = 0f;
+  var li = global_id.z * e.size.x * e.size.y + global_id.x;
+  var ri = global_id.z * dlds.size.x * dlds.size.y + global_id.y;
+  for (var i = 0u; i < e.size.y; i++) {
+    s += e.data[li] * dlds.data[ri];
+    li += e.size.x;
+    ri += dlds.size.x;
+  }
+  gtattw.data[global_id.z * gtattw.size.y * gtattw.size.x + global_id.y * gtattw.size.x + global_id.x] = s;
+}

shader/gradtout.wgsl

@@ -0,0 +1,28 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> a: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> dldo: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> gtout: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  gtout.size = a.size;
+  gtout.size.y = dldo.size.x;
+
+  var s = 0f;
+  var ai = global_id.z * a.size.y * a.size.x + global_id.x;
+  var di = global_id.z * dldo.size.x * dldo.size.y + global_id.y;
+  for (var i = 0u; i < a.size.y; i++) {
+    s += a.data[ai] * dldo.data[di];
+    ai += a.size.x;
+    di += dldo.size.x;
+  }
+
+  gtout.data[global_id.z * gtout.size.y * gtout.size.x + global_id.y * gtout.size.x + global_id.x] = s;
+}

shader/hotones.wgsl

@@ -0,0 +1,19 @@
+struct UVectorBatch {
+  size: vec4<u32>,
+  data: array<u32>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> v: UVectorBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.data[global_id.z * result.size.y * result.size.x + global_id.y * result.size.x + global_id.x] = select(0f, 1f, global_id.x == v.data[global_id.z * v.size.x + global_id.y]);
+}

shader/isum.wgsl

@@ -0,0 +1,37 @@
+alias number = f32;
+
+struct VectorBatch {
+  size: vec4<u32>,
+  data: array<u32>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> v: VectorBatch;
+@group(0) @binding(1)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = m.size;
+
+  var vc = result.size.x / 4;
+  var resi = (global_id.z * result.size.y + global_id.y) * vc + global_id.x;
+  result.data[resi] = vec4(0f);
+
+  var mi = global_id.z * m.size.y * vc + global_id.x;
+  var vi = global_id.z * v.size.x;
+  for (var i = 0u; i < v.size.x; i++) {
+    if (global_id.y == v.data[vi]) {
+      result.data[resi] += m.data[mi];
+    }
+    mi += vc;
+    vi++;
+  }
+}

shader/mul.wgsl

@@ -0,0 +1,31 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> r: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size.x = r.size.x;
+  result.size.y = l.size.y;
+  result.size.z = l.size.z;
+
+  var s = 0f;
+  var li = global_id.z * l.size.x * l.size.y + global_id.y * l.size.x;
+  var ri = global_id.z * r.size.x * r.size.y + global_id.x;
+  for (var i = 0u; i < l.size.x; i++) {
+    s += l.data[li] * r.data[ri];
+    li++;
+    ri += r.size.x;
+  }
+
+  result.data[global_id.z * result.size.y * result.size.x + global_id.y * result.size.x + global_id.x] = s;
+}

shader/otdldo.wgsl

@@ -0,0 +1,29 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> dldo: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> tout: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> otdldo: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  otdldo.size.x = dldo.size.y;
+  otdldo.size.y = tout.size.x;
+  otdldo.size.z = dldo.size.z;
+
+  var s = 0f;
+  var li = (global_id.z * dldo.size.y + global_id.x) * dldo.size.x;
+  var ri = global_id.z * tout.size.x * tout.size.y + global_id.y;
+  for (var i = 0u; i < dldo.size.x; i++) {
+    s += dldo.data[li] * tout.data[ri];
+    li++;
+    ri += tout.size.x;
+  }
+
+  otdldo.data[(global_id.z * otdldo.size.y + global_id.y) * otdldo.size.x + global_id.x] = s;
+}

shader/samples.wgsl

@@ -0,0 +1,39 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+struct FVectorBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+struct UVectorBatch {
+  size: vec4<u32>,
+  data: array<u32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> r: FVectorBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: UVectorBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size.x = m.size.y;
+  result.size.z = m.size.z;
+  var resi = global_id.z * result.size.x + global_id.x;
+  var tp = r.data[global_id.z * r.size.x + global_id.x];
+  var p = 0f;
+  var mi = (global_id.z * m.size.y + global_id.x) * m.size.x;
+  for (var i = 0u; i < m.size.x; i++) {
+    p += m.data[mi];
+    if (p >= tp) {
+      result.data[resi] = i;
+      break;
+    }
+    mi++;
+  }
+}

shader/softmaxexp.wgsl

@@ -0,0 +1,25 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> b: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = b.size;
+
+  var vc = result.size.x / 4;
+  var resi = (global_id.z * result.size.y + global_id.y) * vc + global_id.x;
+  var mv = b.data[resi];
+  var bi = (global_id.z * b.size.y + global_id.y) * vc;
+  for (var i = 0u; i < vc; i++) {
+    mv = max(mv, b.data[bi]);
+    bi++;
+  }
+  var ms = max(max(mv.x, mv.y), max(mv.z, mv.w));
+  result.data[resi] = exp(b.data[resi] - vec4(ms));
+}

shader/softmaxnorm.wgsl

@@ -0,0 +1,25 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> b: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = b.size;
+
+  var vc = result.size.x / 4;
+  var resi = (global_id.z * result.size.y + global_id.y) * vc + global_id.x;
+  var sv = vec4(0f);
+  var bi = (global_id.z * b.size.y + global_id.y) * vc;
+  for (var i = 0u; i < vc; i++) {
+    sv += b.data[bi];
+    bi++;
+  }
+  var ss = sv.x + sv.y + sv.z + sv.w;
+  result.data[resi] = b.data[resi] / vec4(ss);
+}

shader/tot.wgsl

@@ -0,0 +1,32 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+struct Matrix {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: Matrix;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = m.size;
+  result.size.z = 1;
+
+  var vc = result.size.x / 4;
+  var resi = global_id.y * vc + global_id.x;
+  result.data[resi] = vec4(0f);
+
+  var mi = resi;
+  for (var z = 0u; z < m.size.z; z++) {
+    result.data[resi] += m.data[mi];
+    mi += m.size.y * vc;
+  }
+}

shader/totdldowa.wgsl

@@ -0,0 +1,29 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+@group(0) @binding(0)
+var<storage, read> wa: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> otdldo: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> totdldowa: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  totdldowa.size.x = otdldo.size.y;
+  totdldowa.size.y = wa.size.x;
+  totdldowa.size.z = otdldo.size.z;
+
+  var s = 0f;
+  var li = global_id.z * otdldo.size.x * otdldo.size.y + global_id.x * otdldo.size.x;
+  var ri = global_id.z * wa.size.x * wa.size.y + global_id.y;
+  for (var i = 0u; i < otdldo.size.x; i++) {
+    s += otdldo.data[li] * wa.data[ri];
+    li++;
+    ri += wa.size.x;
+  }
+
+  totdldowa.data[(global_id.z * totdldowa.size.y + global_id.y) * totdldowa.size.x + global_id.x] = s;
+}

shader/transpose.wgsl

@@ -0,0 +1,20 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+@group(0) @binding(0)
+var<storage, read> in: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> out: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  out.size.x = in.size.y;
+  out.size.y = in.size.x;
+  out.size.z = in.size.z;
+  var l = global_id.z * out.size.y * out.size.x;
+  out.data[l + global_id.y * out.size.x + global_id.x] = in.data[l + global_id.x * in.size.x + global_id.y];
+}

shader/unavg.wgsl

@@ -0,0 +1,26 @@
+alias number = f32;
+
+struct Matrix {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> m: Matrix;
+@group(0) @binding(1)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>, @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+  result.size = m.size;
+  result.size.z = num_workgroups.z * 2;
+
+  var vc = result.size.x / 4;
+  var mi = global_id.y * vc + global_id.x;
+  result.data[global_id.z * result.size.y * vc + mi] = m.data[mi];
+}

src/data/gpu/gpu-struct.ts

@@ -0,0 +1,301 @@
+import fs from 'node:fs';
+import type { Scalar } from './scalar';
+
+interface ComputePipeline {
+  pipeline: GPUComputePipeline;
+  layout: GPUBindGroupLayout;
+}
+
+const PIPELINES: Map<GPUDevice, Map<string, ComputePipeline>> = new Map();
+
+const COMPWISE_SHADER = {
+  sum: fs.readFileSync('shader/csum.wgsl').toString(),
+  sub: fs.readFileSync('shader/csub.wgsl').toString(),
+  mul: fs.readFileSync('shader/cmul.wgsl').toString(),
+  div: fs.readFileSync('shader/cdiv.wgsl').toString(),
+  scale: fs.readFileSync('shader/cscale.wgsl').toString(),
+};
+
+export function createComputePipeline(
+  device: GPUDevice,
+  name: string,
+  shader: () => string,
+  entry: string = 'main',
+) {
+  if (!PIPELINES.has(device)) {
+    PIPELINES.set(device, new Map());
+  }
+
+  const d = PIPELINES.get(device)!;
+  if (!d.has(name)) {
+    const pipeline = device.createComputePipeline({
+      label: name,
+      layout: 'auto',
+      compute: {
+        module: device.createShaderModule({ code: shader() }),
+        entryPoint: entry,
+      },
+    });
+    const layout = pipeline.getBindGroupLayout(0);
+
+    d.set(name, { pipeline, layout });
+  }
+
+  return d.get(name)!;
+}
+
+export abstract class GPUStruct<A extends Uint32Array | Int32Array | Float32Array> {
+  public static dunit = 8;
+  public static lunit = 2;
+
+  protected readonly _label: string;
+  protected readonly _device: GPUDevice;
+  protected readonly _array: new (size: number | ArrayBufferLike | ArrayLike<number>) => A;
+  protected readonly _inbuffer: GPUBuffer;
+  protected readonly _dummybuffer: GPUBuffer;
+  protected readonly _outbuffer: GPUBuffer;
+  protected _op: 'reading' | 'writing' | null = null;
+  protected _dirty: 'front' | 'back' | null = 'front';
+
+  protected abstract workgroups(pm?: number): [number, number, number];
+  protected abstract workgroup(): [number, number, number];
+
+  public readonly size: Int32Array;
+  public readonly data: A;
+
+  public get type(): 'u32' | 'i32' | 'f32' {
+    switch (this._array as unknown) {
+      case Uint32Array: return 'u32';
+      case Int32Array: return 'i32';
+      case Float32Array: return 'f32';
+      default: throw new Error('Unknown array type.');
+    }
+  }
+
+  public get length(): number {
+    return this.data.length;
+  }
+
+  public get byteLength(): number {
+    return this.size.byteLength + this.data.byteLength;
+  }
+
+  public constructor(
+    device: GPUDevice,
+    array: new (size: number) => A,
+    label: string,
+    size: [number, number, number, number],
+  ) {
+    this._device = device;
+    this._array = array;
+    this._label = label ?? 'struct';
+    this.size = new Int32Array(size);
+    this.data = new array(size[0] * size[1] * size[2] * size[3]);
+
+    this._inbuffer = device.createBuffer({
+      label: `${this._label}in`,
+      size: Math.max(32, this.byteLength),
+      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+    });
+
+    this._dummybuffer = device.createBuffer({
+      label: `${label}dummy`,
+      size: 0,
+      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+    });
+
+    this._outbuffer = this._device.createBuffer({
+      label: `${label}out`,
+      size: Math.max(32, this.byteLength),
+      usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
+    });
+  }
+
+  public write(): void {
+    if (!this._dirty) {
+      return;
+    }
+
+    if (this._op) {
+      throw new Error('Sync in progress.');
+    }
+
+    this._op = 'writing';
+    this._device.queue.writeBuffer(this._inbuffer, 0, new Uint8Array(this.size.buffer));
+    this._device.queue.writeBuffer(this._inbuffer, this.size.byteLength, new Uint8Array(this.data.buffer));
+    this._dirty = null;
+    this._op = null;
+  }
+
+  public async read(): Promise<void> {
+    if (!this._dirty) {
+      return;
+    }
+
+    if (this._op) {
+      throw new Error('Sync in progress.');
+    }
+
+    this._op = 'reading';
+
+    const copy = this._device.createCommandEncoder();
+    copy.copyBufferToBuffer(this._inbuffer, 0, this._outbuffer, 0, this.byteLength);
+    this._device.queue.submit([copy.finish()]);
+
+    await this._device.queue.onSubmittedWorkDone();
+
+    await this._outbuffer.mapAsync(GPUMapMode.READ);
+    this.data.set(new this._array(this._outbuffer.getMappedRange()).slice(this.size.length));
+    this._outbuffer.unmap();
+    this._dirty = null;
+    this._op = null;
+  }
+
+  public buffer(): GPUBuffer {
+    if (this._dirty === 'front') {
+      this.write();
+    }
+    return this._inbuffer;
+  }
+
+  protected workgroupify(shader: string): string {
+    const [x, y, z] = this.workgroup();
+    if (!x || !y || !z) {
+      throw new Error(`Invalid workgroup size: (${x}, ${y}, ${z}).`);
+    }
+    return shader.replace(/@workgroup_size\(.*?\)/g, `@workgroup_size(${x}, ${y}, ${z})`);
+  }
+
+  protected inplacify(shader: string): string {
+    let r = '';
+    return shader.replace(/\w+_or_result/g, v => r = r ? 'result' : v); // The first occurrence must be the declaration.
+  }
+
+  protected typify(shader: string): string {
+    return shader.replace(/alias number = .+?;/g, `alias number = ${this.type};`);
+  }
+
+  public calculate(
+    name: string,
+    shader: string,
+    args: GPUStruct<Uint32Array | Int32Array | Float32Array>[],
+    ip = false,
+    pm = 1,
+  ): typeof this {
+    const encoder = this._device.createCommandEncoder();
+    const pass = encoder.beginComputePass();
+    const { pipeline, layout } = createComputePipeline(this._device, `${name}${ip ? '_inplace_' : '_'}${this.type}`, () => this.workgroupify(this.typify(ip ? this.inplacify(shader) : shader)));
+
+    pass.setPipeline(pipeline);
+    pass.setBindGroup(0, this._device.createBindGroup({
+      layout, entries: [
+        ...args.map((s, i) => ({ binding: i, resource: { buffer: s.buffer() } })),
+        { binding: args.length + +ip, resource: { buffer: this.buffer() } },
+      ],
+    }));
+
+    const [x, y, z] = this.workgroups(pm);
+    if (!x || !y || !z || Math.floor(x) !== x || Math.floor(y) !== y || Math.floor(z) !== z) {
+      throw new Error(`Invalid workgroup count: (${x}, ${y}, ${z}).`);
+    }
+    pass.dispatchWorkgroups(x, y, z);
+    pass.end();
+
+    this._device.queue.submit([encoder.finish()]);
+
+    this._dirty = 'back';
+
+    return this;
+  }
+
+  protected static assertUnique(...ms: GPUStruct<Uint32Array | Int32Array | Float32Array>[]): void {
+    const map: Map<GPUStruct<Uint32Array | Int32Array | Float32Array>, boolean> = new Map();
+    for (const m of ms) {
+      if (map.has(m)) {
+        throw new Error('Can not use the same GPU structure multiple times.');
+      }
+      map.set(m, true);
+    }
+  }
+
+  protected static assertColocated(...ms: GPUStruct<Uint32Array | Int32Array | Float32Array>[]): void {
+    const map: Map<GPUDevice, boolean> = new Map();
+    for (const m of ms) {
+      map.set(m._device, true);
+      if (map.size > 1) {
+        throw new Error('Can not use GPU structures stored on multiple devices.');
+      }
+    }
+  }
+
+  protected static assertSimilar(...ms: GPUStruct<Uint32Array | Int32Array | Float32Array>[]): void {
+    for (const m of ms) {
+      if (m.length !== ms[0].length
+        || m.workgroup()[0] !== ms[0].workgroup()[0]
+        || m.workgroup()[1] !== ms[0].workgroup()[1]
+        || m.workgroup()[2] !== ms[0].workgroup()[2]) {
+        throw new Error('Can not use GPU structures with different dimensions.');
+      }
+    }
+  }
+
+  protected static assertPackable(pm: number, ...ms: GPUStruct<Uint32Array | Int32Array | Float32Array>[]): void {
+    for (const m of ms) {
+      if (m.workgroups(pm).some(v => v < 1)) {
+        throw new Error(`Struct ${m._label} size must be divisible by ${pm}.`);
+      }
+    }
+  }
+
+  public static compwise<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(op: keyof typeof COMPWISE_SHADER, l: S, r: S, result: S): S {
+    this.assertUnique(l, r);
+    this.assertUnique(l, result);
+    this.assertColocated(l, r, result);
+    this.assertPackable(4, l, r, result);
+
+    return result.calculate(op, COMPWISE_SHADER[op], result === r ? [l] : [l, r], result === r, 4);
+  }
+
+  public static csum<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(l: S, r: S, result: S): S {
+    return GPUStruct.compwise('sum', l, r, result);
+  }
+
+  public static csub<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(l: S, r: S, result: S): S {
+    return GPUStruct.compwise('sub', l, r, result);
+  }
+
+  public static cmul<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(l: S, r: S, result: S): S {
+    return GPUStruct.compwise('mul', l, r, result);
+  }
+
+  public static cdiv<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(l: S, r: S, result: S): S {
+    return GPUStruct.compwise('div', l, r, result);
+  }
+
+  public static cscale<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(l: Scalar<A>, r: S, result: S): S {
+    this.assertUnique(l, r);
+    this.assertUnique(l, result);
+    this.assertColocated(l, r, result);
+    this.assertPackable(4, r, result);
+
+    return result.calculate('scale', COMPWISE_SHADER.scale, result === r ? [l] : [l, r], result === r, 4);
+  }
+}

src/data/gpu/matrix-batch.ts

@@ -0,0 +1,289 @@
+import fs from 'node:fs';
+import { Matrix as CPUMatrix, Vector as CPUVector } from '../math';
+import { GPUStruct } from './gpu-struct';
+import { Matrix } from './matrix';
+import { ScalarBatch } from './scalar-batch';
+import { VectorBatch } from './vector-batch';
+
+const SHADER_SCALE = fs.readFileSync('shader/bscale.wgsl').toString();
+const SHADER_ISUM = fs.readFileSync('shader/isum.wgsl').toString();
+const SHADER_MUL = fs.readFileSync('shader/mul.wgsl').toString();
+const SHADER_TRANS = fs.readFileSync('shader/transpose.wgsl').toString();
+const SHADER_AVG = fs.readFileSync('shader/avg.wgsl').toString();
+const SHADER_TOT = fs.readFileSync('shader/tot.wgsl').toString();
+
+export class MatrixBatch<
+  L extends number,
+  H extends number,
+  W extends number,
+  A extends Uint32Array | Int32Array | Float32Array,
+  T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+> extends GPUStruct<A> {
+  public get width(): W {
+    return this.size[0] as W;
+  }
+
+  public get height(): H {
+    return this.size[1] as H;
+  }
+
+  public get layers(): L {
+    return this.size[2] as L;
+  }
+
+  protected override workgroups(pm: number = 1): [number, number, number] {
+    return [this.width / GPUStruct.dunit / pm, this.height / GPUStruct.dunit, this.layers / MatrixBatch.lunit];
+  }
+
+  protected override workgroup(): [number, number, number] {
+    return [GPUStruct.dunit, GPUStruct.dunit, MatrixBatch.lunit];
+  }
+
+  public constructor(device: GPUDevice, array: new (size: number) => A, data: CPUVector<L, CPUMatrix<H, W, T>>);
+  public constructor(device: GPUDevice, array: new (size: number) => A, l: L, h: H, w: W, data?: CPUVector<L, CPUMatrix<H, W, T>>);
+  public constructor(device: GPUDevice, array: new (size: number) => A, ...args: unknown[]) {
+    let [l, h, w]: [number, number, number] = [1, 1, 1];
+    if (typeof args[0] === 'number') {
+      [l, h, w] = args as [L, H, W];
+    } else {
+      const [data] = args as [CPUVector<L, CPUMatrix<H, W, T>>];
+      [l, h, w] = [data.length, data[0].length, data[0][0].length, 1];
+    }
+
+    super(device, array, `${h}x${w}[${l}]`, [w, h, l, 1]);
+
+    MatrixBatch.assertPackable(1, this);
+
+    const data = (typeof args[0] === 'number' ? args[3] : args[0]) as CPUVector<L, CPUMatrix<H, W, T>> | undefined;
+    if (data) {
+      this.set(data);
+    }
+  }
+
+  public index(l: number, y: number, x: number): number {
+    return (l * this.height + y) * this.width + x;
+  }
+
+  public async set(y: number, x: number, v: T): Promise<void>;
+  public async set(data: CPUVector<L, CPUMatrix<H, W, T>>): Promise<void>;
+  public async set(slice: T[][][], ob?: number, oy?: number, ox?: number): Promise<void>;
+  public async set(m: MatrixBatch<L, H, W, A, T>): Promise<void>;
+  public async set(...args: unknown[]): Promise<void> {
+    if (typeof args[0] === 'number') {
+      const [l, y, x, v]: number[] = args as number[];
+
+      if (this._dirty === 'back') {
+        await this.read();
+      }
+
+      this.data[this.index(l, y, x)] = v;
+      this._dirty = 'front';
+    } else if (Array.isArray(args[0])) {
+      const [data, ol = 0, oy = 0, ox = 0] = args as [T[][][], number | undefined, number | undefined, number | undefined];
+
+      if (this._dirty === 'back' && (ol || oy || ox || data.length < this.layers || data[0].length < this.height || data[0][0].length < this.width)) {
+        await this.read();
+      }
+
+      let l = ol;
+      let y = oy;
+      let x = ox;
+      for (const mat of data ?? []) {
+        if (l >= this.layers) {
+          break;
+        }
+        for (const row of mat ?? []) {
+          if (y >= this.height) {
+            break;
+          }
+          for (const v of row) {
+            if (x >= this.width) {
+              break;
+            }
+            this.data[this.index(l, y, x)] = v;
+            x++;
+          }
+          x = ox;
+          y++;
+        }
+        y = oy;
+        l++;
+      }
+      this._dirty = 'front';
+    } else {
+      const [batch] = args as [MatrixBatch<L, H, W, A, T>];
+      const copy = this._device.createCommandEncoder();
+      copy.copyBufferToBuffer(batch.buffer(), 0, this.buffer(), 0, this.byteLength);
+      this._device.queue.submit([copy.finish()]);
+      this._dirty = 'back';
+    }
+  }
+
+  public async get(l: number, y: number, x: number): Promise<T>;
+  public async get(l: number): Promise<CPUMatrix<H, W, T>>;
+  public async get(): Promise<CPUVector<L, CPUMatrix<H, W, T>>>;
+  public async get(l?: number, y?: number, x?: number): Promise<T | CPUMatrix<H, W, T> | CPUVector<L, CPUMatrix<H, W, T>>> {
+    if (this._dirty === 'back') {
+      await this.read();
+    }
+    if (typeof x === 'number') {
+      return this.data[this.index(l!, y!, x!)] as T;
+    } else if (typeof l === 'number') {
+      const result: CPUMatrix<H, W, T> = [] as unknown as CPUMatrix<H, W, T>;
+      for (let y: number = 0; y < this.height; y++) {
+        const row: CPUVector<W, T> = [] as unknown as CPUVector<W, T>;
+        for (let x: number = 0; x < this.width; x++) {
+          row.push(this.data[this.index(l, y, x)] as T);
+        }
+        result.push(row);
+      }
+      return result;
+    } else {
+      const result: CPUVector<L, CPUMatrix<H, W, T>> = [] as unknown as CPUVector<L, CPUMatrix<H, W, T>>;
+      for (let l: number = 0; l < this.layers; l++) {
+        const layer: CPUMatrix<H, W, T> = [] as unknown as CPUMatrix<H, W, T>;
+        for (let y: number = 0; y < this.height; y++) {
+          const row: CPUVector<W, T> = [] as unknown as CPUVector<W, T>;
+          for (let x: number = 0; x < this.width; x++) {
+            row.push(this.data[this.index(l, y, x)] as T);
+          }
+          layer.push(row);
+        }
+        result.push(layer);
+      }
+      return result;
+    }
+  }
+
+  protected static override assertPackable(pm: number, ...ms: MatrixBatch<number, number, number, Uint32Array | Int32Array | Float32Array>[]): void {
+    for (const m of ms) {
+      if (m.layers % MatrixBatch.lunit !== 0 || m.height % MatrixBatch.dunit !== 0 || m.width % (MatrixBatch.dunit * pm) !== 0) {
+        throw new Error(`Batch size ${m.height}x${m.width}[${m.layers}] must be divisible by ${MatrixBatch.dunit}x${MatrixBatch.dunit * pm}[${MatrixBatch.lunit}].`);
+      }
+    }
+  }
+
+  public static sum<
+    L extends number,
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: MatrixBatch<L, H, W, A, T>, r: MatrixBatch<L, H, W, A, T>, result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, l.height, l.width)): MatrixBatch<L, H, W, A, T> {
+    return GPUStruct.csum(l, r, result);
+  }
+
+  public static sub<
+    L extends number,
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: MatrixBatch<L, H, W, A, T>, r: MatrixBatch<L, H, W, A, T>, result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, l.height, l.width)): MatrixBatch<L, H, W, A, T> {
+    return GPUStruct.csub(l, r, result);
+  }
+
+  public static isum<
+    L extends number,
+    H extends number,
+    I extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    lis: VectorBatch<L, I, Int32Array>,
+    m: MatrixBatch<L, H, W, A, T>,
+    result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(m._device, m._array, m.layers, m.height, m.width),
+  ): MatrixBatch<L, H, W, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m, result);
+
+    return result.calculate('isum', SHADER_ISUM, [lis, m], false, 4);
+  }
+
+  public static scale<
+    L extends number,
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: ScalarBatch<L, A, T>, r: MatrixBatch<L, H, W, A, T>, result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(r._device, r._array, r.layers, r.height, r.width)): MatrixBatch<L, H, W, A, T> {
+    this.assertUnique(l, r);
+    this.assertUnique(l, result);
+    this.assertColocated(l, r, result);
+    this.assertPackable(4, r, result);
+
+    return result.calculate('scale', SHADER_SCALE, result === r ? [l] : [l, r], result === r, 4);
+  }
+
+  public static mul<
+    L extends number,
+    H extends number,
+    S extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: MatrixBatch<L, H, S, A, T>,
+    r: MatrixBatch<L, S, W, A, T>,
+    result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, l.height, r.width),
+  ): MatrixBatch<L, H, W, A, T> {
+    this.assertUnique(l, r, result);
+    this.assertColocated(l, r, result);
+    this.assertPackable(1, l, r, result);
+
+    return result.calculate('mul', SHADER_MUL, [l, r]);
+  }
+
+  public static transpose<
+    L extends number,
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    m: MatrixBatch<L, H, W, A, T>,
+    result: MatrixBatch<L, W, H, A, T> = new MatrixBatch(m._device, m._array, m.layers, m.width, m.height),
+  ): MatrixBatch<L, W, H, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(1, m, result);
+
+    return result.calculate('transpose', SHADER_TRANS, [m]);
+  }
+
+  public static avg<
+    L extends number,
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    m: MatrixBatch<L, H, W, A, T>,
+    result: Matrix<H, W, A, T> = new Matrix(m._device, m._array, m.height, m.width),
+  ): Matrix<H, W, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m);
+
+    return result.calculate('avg', SHADER_AVG, [m], false, 4);
+  }
+
+  public static tot<
+    L extends number,
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    m: MatrixBatch<L, H, W, A, T>,
+    result: Matrix<H, W, A, T> = new Matrix(m._device, m._array, m.height, m.width),
+  ): Matrix<H, W, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m);
+
+    return result.calculate('tot', SHADER_TOT, [m], false, 4);
+  }
+}

src/data/gpu/matrix.ts

@@ -0,0 +1,233 @@
+import fs from 'node:fs';
+import { Matrix as CPUMatrix, Vector as CPUVector } from '../math';
+import { GPUStruct } from './gpu-struct';
+import { MatrixBatch } from './matrix-batch';
+import { Scalar } from './scalar';
+import { Vector } from './vector';
+
+const SHADER_ISUM = fs.readFileSync('shader/isum.wgsl').toString();
+const SHADER_MUL = fs.readFileSync('shader/mul.wgsl').toString();
+const SHADER_TRANS = fs.readFileSync('shader/transpose.wgsl').toString();
+const SHADER_UNAVG = fs.readFileSync('shader/unavg.wgsl').toString();
+
+export class Matrix<
+  H extends number,
+  W extends number,
+  A extends Uint32Array | Int32Array | Float32Array,
+  T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+> extends GPUStruct<A> {
+  public get width(): W {
+    return this.size[0] as W;
+  }
+
+  public get height(): H {
+    return this.size[1] as H;
+  }
+
+  protected override workgroups(pm: number = 1): [number, number, number] {
+    return [this.width / GPUStruct.dunit / pm, this.height / GPUStruct.dunit, 1];
+  }
+
+  protected override workgroup(): [number, number, number] {
+    return [GPUStruct.dunit, GPUStruct.dunit, 1];
+  }
+
+  public constructor(device: GPUDevice, array: new (size: number) => A, data: CPUMatrix<H, W, T>);
+  public constructor(device: GPUDevice, array: new (size: number) => A, h: H, w: W, data?: CPUMatrix<H, W, T>);
+  public constructor(device: GPUDevice, array: new (size: number) => A, ...args: unknown[]) {
+    let [h, w]: [number, number] = [1, 1];
+    if (typeof args[0] === 'number') {
+      [h, w] = args as [H, W];
+    } else {
+      const [data] = args as [CPUMatrix<H, W, T>];
+      [h, w] = [data.length, data[0].length];
+    }
+
+    super(device, array, `${h}x${w}`, [w, h, 1, 1]);
+
+    Matrix.assertPackable(1, this);
+
+    const data = (typeof args[0] === 'number' ? args[2] : args[0]) as CPUMatrix<H, W, T> | undefined;
+    if (data) {
+      this.set(data);
+    }
+  }
+
+  public index(y: number, x: number): number {
+    return y * this.width + x;
+  }
+
+  public async set(y: number, x: number, v: T): Promise<void>;
+  public async set(data: CPUMatrix<H, W, T>): Promise<void>;
+  public async set(slice: T[][], oy?: number, ox?: number): Promise<void>;
+  public async set(m: Matrix<H, W, A, T>): Promise<void>;
+  public async set(...args: unknown[]): Promise<void> {
+    if (typeof args[0] === 'number') {
+      const [y, x, v]: number[] = args as number[];
+
+      if (this._dirty === 'back') {
+        await this.read();
+      }
+
+      this.data[this.index(y, x)] = v;
+      this._dirty = 'front';
+    } else if (Array.isArray(args[0])) {
+      const [data, oy = 0, ox = 0] = args as [T[][], number | undefined, number | undefined];
+
+      if (this._dirty === 'back' && (oy || ox || data.length < this.height || data[0].length < this.width)) {
+        await this.read();
+      }
+
+      let y = oy;
+      let x = ox;
+      for (const row of data ?? []) {
+        if (y >= this.height) {
+          break;
+        }
+        for (const v of row) {
+          if (x >= this.width) {
+            break;
+          }
+          this.data[this.index(y, x)] = v;
+          x++;
+        }
+        x = ox;
+        y++;
+      }
+      this._dirty = 'front';
+    } else {
+      const [matrix] = args as [Matrix<H, W, A, T>];
+      const copy = this._device.createCommandEncoder();
+      copy.copyBufferToBuffer(matrix.buffer(), 0, this.buffer(), 0, this.byteLength);
+      this._device.queue.submit([copy.finish()]);
+      this._dirty = 'back';
+    }
+  }
+
+  public async get(y: number, x: number): Promise<T>;
+  public async get(): Promise<CPUMatrix<H, W, T>>;
+  public async get(y?: number, x?: number): Promise<T | CPUMatrix<H, W, T>> {
+    if (this._dirty === 'back') {
+      await this.read();
+    }
+    if (typeof x === 'number') {
+      return this.data[this.index(y!, x!)] as T;
+    } else {
+      const result: CPUMatrix<H, W, T> = [] as unknown as CPUMatrix<H, W, T>;
+      for (let y: number = 0; y < this.height; y++) {
+        const row: CPUVector<W, T> = [] as unknown as CPUVector<W, T>;
+        for (let x: number = 0; x < this.width; x++) {
+          row.push(this.data[this.index(y, x)] as T);
+        }
+        result.push(row);
+      }
+      return result;
+    }
+  }
+
+  protected static override assertPackable(pm: number, ...ms: Matrix<number, number, Uint32Array | Int32Array | Float32Array>[]): void {
+    for (const m of ms) {
+      if (m.height % Matrix.dunit !== 0 || m.width % (Matrix.dunit * pm) !== 0) {
+        throw new Error(`Matrix size ${m.height}x${m.width} must be divisible by ${Matrix.dunit}x${Matrix.dunit * pm}.`);
+      }
+    }
+  }
+
+  public static sum<
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Matrix<H, W, A, T>, r: Matrix<H, W, A, T>, result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.height, l.width)): Matrix<H, W, A, T> {
+    return GPUStruct.csum(l, r, result);
+  }
+
+  public static sub<
+    H extends number,
+    W extends number,
+
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Matrix<H, W, A, T>, r: Matrix<H, W, A, T>, result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.height, l.width)): Matrix<H, W, A, T> {
+    return GPUStruct.csub(l, r, result);
+  }
+
+  public static isum<
+    H extends number,
+    I extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    lis: Vector<I, Int32Array>,
+    m: Matrix<H, W, A, T>,
+    result: Matrix<H, W, A, T> = new Matrix(m._device, m._array, m.height, m.width),
+  ): Matrix<H, W, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m, result);
+
+    return result.calculate('isum', SHADER_ISUM, [lis, m], false, 4);
+  }
+
+  public static scale<
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Scalar<A, T>, r: Matrix<H, W, A, T>, result: Matrix<H, W, A, T> = new Matrix(r._device, r._array, r.height, r.width)): Matrix<H, W, A, T> {
+    return GPUStruct.cscale(l, r, result);
+  }
+
+  public static mul<
+    H extends number,
+    S extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: Matrix<H, S, A, T>,
+    r: Matrix<S, W, A, T>,
+    result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.height, r.width),
+  ): Matrix<H, W, A, T> {
+    this.assertUnique(l, r, result);
+    this.assertColocated(l, r, result);
+    this.assertPackable(1, l, r, result);
+
+    return result.calculate('mul', SHADER_MUL, [l, r]);
+  }
+
+  public static transpose<
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    m: Matrix<H, W, A, T>,
+    result: Matrix<W, H, A, T> = new Matrix(m._device, m._array, m.width, m.height),
+  ): Matrix<W, H, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(1, m, result);
+
+    return result.calculate('transpose', SHADER_TRANS, [m]);
+  }
+
+  public static unavg<
+    L extends number,
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: L,
+    m: Matrix<H, W, A, T>,
+    result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(m._device, m._array, l, m.height, m.width),
+  ): MatrixBatch<L, H, W, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m);
+
+    return result.calculate('unavg', SHADER_UNAVG, [m], false, 4);
+  }
+}

src/data/gpu/scalar-batch.ts

@@ -0,0 +1,176 @@
+import fs from 'node:fs';
+import { Vector as CPUVector } from '../math';
+import { GPUStruct } from './gpu-struct';
+import { Scalar } from './scalar';
+
+const SHADER_AVG = fs.readFileSync('shader/avg.wgsl').toString();
+const SHADER_TOT = fs.readFileSync('shader/tot.wgsl').toString();
+
+export class ScalarBatch<
+  L extends number,
+  A extends Uint32Array | Int32Array | Float32Array,
+  T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+> extends GPUStruct<A> {
+  public get layers(): L {
+    return this.size[2] as L;
+  }
+
+  protected override workgroups(): [number, number, number] {
+    return [1, 1, this.layers / GPUStruct.lunit];
+  }
+
+  protected override workgroup(): [number, number, number] {
+    return [1, 1, GPUStruct.lunit];
+  }
+
+  public constructor(device: GPUDevice, array: new (size: number) => A, data: CPUVector<L, T>);
+  public constructor(device: GPUDevice, array: new (size: number) => A, l: L, data?: CPUVector<L, T>);
+  public constructor(device: GPUDevice, array: new (size: number) => A, ...args: unknown[]) {
+    const l: number = typeof args[0] === 'number'
+      ? args[0] as L
+      : (args[0] as CPUVector<L, T>).length;
+
+    super(device, array, `[${l}]`, [1, 1, l, 1]);
+
+    ScalarBatch.assertPackable(1, this);
+
+    const data = (typeof args[0] === 'number' ? args[1] : args[0]) as CPUVector<L, T> | undefined;
+    if (data) {
+      this.set(data);
+    }
+  }
+
+  public async set(l: number, v: T): Promise<void>;
+  public async set(data: CPUVector<L, T>): Promise<void>;
+  public async set(slice: T[], ol?: number): Promise<void>;
+  public async set(b: ScalarBatch<L, A, T>): Promise<void>;
+  public async set(...args: unknown[]): Promise<void> {
+    if (typeof args[0] === 'number') {
+      const [l, v]: number[] = args as number[];
+
+      if (this._dirty === 'back') {
+        await this.read();
+      }
+
+      this.data[l] = v;
+      this._dirty = 'front';
+    } else if (Array.isArray(args[0])) {
+      const [data, ol = 0] = args as [T[], number | undefined];
+
+      if (this._dirty === 'back' && (ol || data.length < this.layers)) {
+        await this.read();
+      }
+
+      let l = ol;
+      for (const v of data ?? []) {
+        if (l >= this.layers) {
+          break;
+        }
+        this.data[l] = v;
+        l++;
+      }
+      this._dirty = 'front';
+    } else {
+      const [batch] = args as [ScalarBatch<L, A, T>];
+      const copy = this._device.createCommandEncoder();
+      copy.copyBufferToBuffer(batch.buffer(), 0, this.buffer(), 0, this.byteLength);
+      this._device.queue.submit([copy.finish()]);
+      this._dirty = 'back';
+    }
+  }
+
+  public async get(l: number): Promise<T>;
+  public async get(): Promise<CPUVector<L, T>>;
+  public async get(l?: number): Promise<T | CPUVector<L, T>> {
+    if (this._dirty === 'back') {
+      await this.read();
+    }
+    if (typeof l === 'number') {
+      return this.data[l] as T;
+    } else {
+      const result: CPUVector<L, T> = [] as unknown as CPUVector<L, T>;
+      for (let l: number = 0; l < this.layers; l++) {
+        result.push(this.data[l] as T);
+      }
+      return result;
+    }
+  }
+
+  protected static override assertPackable(_: number, ...ms: ScalarBatch<number, Uint32Array | Int32Array | Float32Array>[]): void {
+    for (const m of ms) {
+      if (m.layers % ScalarBatch.dunit !== 0) {
+        throw new Error(`Batch size [${m.layers}] must be divisible by [${ScalarBatch.dunit}].`);
+      }
+    }
+  }
+
+  public static sum<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: ScalarBatch<L, A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> {
+    return GPUStruct.csum(l, r, result);
+  }
+
+  public static sub<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: ScalarBatch<L, A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> {
+    return GPUStruct.csub(l, r, result);
+  }
+
+  public static mul<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: ScalarBatch<L, A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> {
+    return GPUStruct.cmul(l, r, result);
+  }
+
+  public static div<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: ScalarBatch<L, A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> {
+    return GPUStruct.cdiv(l, r, result);
+  }
+
+  public static scale<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Scalar<A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(r._device, r._array, r.layers)): ScalarBatch<L, A, T> {
+    return GPUStruct.cscale(l, r, result);
+  }
+
+  public static avg<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    m: ScalarBatch<L, A, T>,
+    result: Scalar<A, T> = new Scalar(m._device, m._array),
+  ): Scalar<A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m);
+
+    return result.calculate('avg', SHADER_AVG, [m], false, 4);
+  }
+
+  public static tot<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    m: ScalarBatch<L, A, T>,
+    result: Scalar<A, T> = new Scalar(m._device, m._array),
+  ): Scalar<A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m);
+
+    return result.calculate('tot', SHADER_TOT, [m], false, 4);
+  }
+}

src/data/gpu/scalar.ts

@@ -0,0 +1,90 @@
+import fs from 'node:fs';
+import { GPUStruct } from './gpu-struct';
+import { ScalarBatch } from './scalar-batch';
+
+const SHADER_UNAVG = fs.readFileSync('shader/unavg.wgsl').toString();
+
+export class Scalar<
+  A extends Uint32Array | Int32Array | Float32Array,
+  T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+> extends GPUStruct<A> {
+  protected override workgroups(): [number, number, number] {
+    return [1, 1, 1];
+  }
+
+  protected override workgroup(): [number, number, number] {
+    return [1, 1, 1];
+  }
+
+  public constructor(device: GPUDevice, array: new (size: number) => A, data?: T) {
+    super(device, array, 'scalar', [1, 1, 1, 1]);
+    this.set(data || 0 as T);
+  }
+
+  public async set(v: T | Scalar<A, T>): Promise<void> {
+    if (typeof v === 'number') {
+      if (this._dirty === 'back') {
+        await this.read();
+      }
+
+      this.data[0] = v;
+      this._dirty = 'front';
+    } else {
+      const copy = this._device.createCommandEncoder();
+      copy.copyBufferToBuffer(v.buffer(), 0, this.buffer(), 0, this.byteLength);
+      this._device.queue.submit([copy.finish()]);
+      this._dirty = 'back';
+    }
+  }
+
+  public async get(): Promise<T> {
+    if (this._dirty === 'back') {
+      await this.read();
+    }
+    return this.data[0] as T;
+  }
+
+  public static sum<
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Scalar<A, T>, r: Scalar<A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> {
+    return GPUStruct.csum(l, r, result);
+  }
+
+  public static sub<
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Scalar<A, T>, r: Scalar<A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> {
+    return GPUStruct.csub(l, r, result);
+  }
+
+  public static mul<
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Scalar<A, T>, r: Scalar<A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> {
+    return GPUStruct.cmul(l, r, result);
+  }
+
+  public static div<
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Scalar<A, T>, r: Scalar<A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> {
+    return GPUStruct.cdiv(l, r, result);
+  }
+
+  public static unavg<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: L,
+    m: Scalar<A, T>,
+    result: ScalarBatch<L, A, T> = new ScalarBatch(m._device, m._array, l),
+  ): ScalarBatch<L, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m);
+
+    return result.calculate('unavg', SHADER_UNAVG, [m], false, 4);
+  }
+}

src/data/gpu/vector-batch.ts

@@ -0,0 +1,221 @@
+import fs from 'node:fs';
+import { Vector as CPUVector } from '../math';
+import { GPUStruct } from './gpu-struct';
+import { ScalarBatch } from './scalar-batch';
+import { Vector } from './vector';
+
+const SHADER_AVG = fs.readFileSync('shader/avg.wgsl').toString();
+const SHADER_TOT = fs.readFileSync('shader/tot.wgsl').toString();
+const SHADER_SCALE = fs.readFileSync('shader/bscale.wgsl').toString();
+const SHADER_DOT = fs.readFileSync('shader/mul.wgsl').toString();
+
+export class VectorBatch<
+  L extends number,
+  D extends number,
+  A extends Uint32Array | Int32Array | Float32Array,
+  T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+> extends GPUStruct<A> {
+  public get dims(): D {
+    return this.size[0] as D;
+  }
+
+  public get layers(): L {
+    return this.size[2] as L;
+  }
+
+  protected override workgroups(pm: number = 1): [number, number, number] {
+    return [this.dims / GPUStruct.dunit / pm, 1, this.layers / GPUStruct.lunit];
+  }
+
+  protected override workgroup(): [number, number, number] {
+    return [GPUStruct.dunit, 1, GPUStruct.lunit];
+  }
+
+  public constructor(device: GPUDevice, array: new (size: number) => A, data: CPUVector<L, CPUVector<D, T>>);
+  public constructor(device: GPUDevice, array: new (size: number) => A, l: L, d: D, data?: CPUVector<L, CPUVector<D, T>>);
+  public constructor(device: GPUDevice, array: new (size: number) => A, ...args: unknown[]) {
+    let [l, d]: [number, number] = [1, 1];
+    if (typeof args[0] === 'number') {
+      [l, d] = args as [L, D];
+    } else {
+      const [data] = args as [CPUVector<L, CPUVector<D, T>>];
+      [l, d] = [data.length, data[0].length];
+    }
+
+    super(device, array, `${d}[${l}]`, [d, 1, l, 1]);
+
+    VectorBatch.assertPackable(1, this);
+
+    const data = (typeof args[0] === 'number' ? args[2] : args[0]) as CPUVector<L, CPUVector<D, T>> | undefined;
+    if (data) {
+      this.set(data);
+    }
+  }
+
+  public index(l: number, i: number): number {
+    return l * this.dims + i;
+  }
+
+  public async set(l: number, i: number, v: T): Promise<void>;
+  public async set(data: CPUVector<L, CPUVector<D, T>>): Promise<void>;
+  public async set(slice: T[][], ol?: number, oi?: number): Promise<void>;
+  public async set(b: VectorBatch<L, D, A, T>): Promise<void>;
+  public async set(...args: unknown[]): Promise<void> {
+    if (typeof args[0] === 'number') {
+      const [l, i, v]: number[] = args as number[];
+
+      if (this._dirty === 'back') {
+        await this.read();
+      }
+
+      this.data[this.index(l, i)] = v;
+      this._dirty = 'front';
+    } else if (Array.isArray(args[0])) {
+      const [data, ol = 0, oi = 0] = args as [T[][], number | undefined, number | undefined];
+
+      if (this._dirty === 'back' && (ol || oi || data.length < this.layers || data[0].length < this.dims)) {
+        await this.read();
+      }
+
+      let l = ol;
+      let i = oi;
+      for (const vec of data ?? []) {
+        if (l >= this.layers) {
+          break;
+        }
+        for (const v of vec) {
+          if (i >= this.dims) {
+            break;
+          }
+          this.data[this.index(l, i)] = v;
+          i++;
+        }
+        i = oi;
+        l++;
+      }
+      this._dirty = 'front';
+    } else {
+      const [batch] = args as [VectorBatch<L, D, A, T>];
+      const copy = this._device.createCommandEncoder();
+      copy.copyBufferToBuffer(batch.buffer(), 0, this.buffer(), 0, this.byteLength);
+      this._device.queue.submit([copy.finish()]);
+      this._dirty = 'back';
+    }
+  }
+
+  public async get(l: number, i: number): Promise<T>;
+  public async get(l: number): Promise<CPUVector<D, T>>;
+  public async get(): Promise<CPUVector<L, CPUVector<D, T>>>;
+  public async get(l?: number, i?: number): Promise<T | CPUVector<D, T> | CPUVector<L, CPUVector<D, T>>> {
+    if (this._dirty === 'back') {
+      await this.read();
+    }
+    if (typeof i === 'number') {
+      return this.data[this.index(l!, i!)] as T;
+    } else if (typeof l === 'number') {
+      const result: CPUVector<D, T> = [] as unknown as CPUVector<D, T>;
+      for (let i: number = 0; i < this.dims; i++) {
+        result.push(this.data[this.index(l, i)] as T);
+      }
+      return result;
+    } else {
+      const result: CPUVector<L, CPUVector<D, T>> = [] as unknown as CPUVector<L, CPUVector<D, T>>;
+      for (let l: number = 0; l < this.layers; l++) {
+        const vec: CPUVector<D, T> = [] as unknown as CPUVector<D, T>;
+        for (let i: number = 0; i < this.dims; i++) {
+          vec.push(this.data[this.index(l, i)] as T);
+        }
+        result.push(vec);
+      }
+      return result;
+    }
+  }
+
+  protected static override assertPackable(pm: number, ...ms: VectorBatch<number, number, Uint32Array | Int32Array | Float32Array>[]): void {
+    for (const m of ms) {
+      if (m.layers % VectorBatch.lunit !== 0 || m.dims % (VectorBatch.dunit * pm) !== 0) {
+        throw new Error(`Batch size ${m.dims}[${m.layers}] must be divisible by ${VectorBatch.dunit * pm}[${VectorBatch.lunit}].`);
+      }
+    }
+  }
+
+  public static sum<
+    L extends number,
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: VectorBatch<L, D, A, T>, r: VectorBatch<L, D, A, T>, result: VectorBatch<L, D, A, T> = new VectorBatch(l._device, l._array, l.layers, l.dims)): VectorBatch<L, D, A, T> {
+    return GPUStruct.csum(l, r, result);
+  }
+
+  public static sub<
+    L extends number,
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: VectorBatch<L, D, A, T>, r: VectorBatch<L, D, A, T>, result: VectorBatch<L, D, A, T> = new VectorBatch(l._device, l._array, l.layers, l.dims)): VectorBatch<L, D, A, T> {
+    return GPUStruct.csub(l, r, result);
+  }
+
+  public static mul<
+    L extends number,
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: ScalarBatch<L, A, T>, r: VectorBatch<L, D, A, T>, result: VectorBatch<L, D, A, T> = new VectorBatch(r._device, r._array, r.layers, r.dims)): VectorBatch<L, D, A, T> {
+    this.assertUnique(l, r);
+    this.assertUnique(l, result);
+    this.assertColocated(l, r, result);
+    this.assertPackable(4, r, result);
+
+    return result.calculate('scale', SHADER_SCALE, result === r ? [l] : [l, r], result === r, 4);
+  }
+
+  public static dot<
+    L extends number,
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: VectorBatch<L, D, A, T>, r: VectorBatch<L, D, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> {
+    throw new Error('Not implemented.');
+
+    this.assertUnique(l, r);
+    this.assertUnique(l, result);
+    this.assertColocated(l, r, result);
+    // this.assertPackable(4, r, result);
+
+    return result.calculate('dot', SHADER_DOT, [l, r]);
+  }
+
+  public static avg<
+    L extends number,
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    m: VectorBatch<L, D, A, T>,
+    result: Vector<D, A, T> = new Vector(m._device, m._array, m.dims),
+  ): Vector<D, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m);
+
+    return result.calculate('avg', SHADER_AVG, [m], false, 4);
+  }
+
+  public static tot<
+    L extends number,
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    m: VectorBatch<L, D, A, T>,
+    result: Vector<D, A, T> = new Vector(m._device, m._array, m.dims),
+  ): Vector<D, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m);
+
+    return result.calculate('tot', SHADER_TOT, [m], false, 4);
+  }
+}

src/data/gpu/vector.ts

@@ -0,0 +1,164 @@
+import fs from 'node:fs';
+import { Vector as CPUVector } from '../math';
+import { GPUStruct } from './gpu-struct';
+import { Scalar } from './scalar';
+import { VectorBatch } from './vector-batch';
+
+const SHADER_UNAVG = fs.readFileSync('shader/unavg.wgsl').toString();
+
+const SHADER_DOT = fs.readFileSync('shader/mul.wgsl').toString();
+
+export class Vector<
+  D extends number,
+  A extends Uint32Array | Int32Array | Float32Array,
+  T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+> extends GPUStruct<A> {
+  public get dims(): D {
+    return this.size[0] as D;
+  }
+
+  protected override workgroups(pm: number = 1): [number, number, number] {
+    return [this.dims / GPUStruct.dunit / pm, 1, 1];
+  }
+
+  protected override workgroup(): [number, number, number] {
+    return [GPUStruct.dunit, 1, 1];
+  }
+
+  public constructor(device: GPUDevice, array: new (size: number) => A, data: CPUVector<D, T>);
+  public constructor(device: GPUDevice, array: new (size: number) => A, d: D, data?: CPUVector<D, T>);
+  public constructor(device: GPUDevice, array: new (size: number) => A, ...args: unknown[]) {
+    const d: number = typeof args[0] === 'number'
+      ? args[0] as D
+      : (args[0] as CPUVector<D, T>).length;
+
+    super(device, array, `${d}`, [d, 1, 1, 1]);
+
+    Vector.assertPackable(1, this);
+
+    const data = (typeof args[0] === 'number' ? args[1] : args[0]) as CPUVector<D, T> | undefined;
+    if (data) {
+      this.set(data);
+    }
+  }
+
+  public async set(i: number, v: T): Promise<void>;
+  public async set(data: CPUVector<D, T>): Promise<void>;
+  public async set(slice: T[], oi?: number): Promise<void>;
+  public async set(b: Vector<D, A, T>): Promise<void>;
+  public async set(...args: unknown[]): Promise<void> {
+    if (typeof args[0] === 'number') {
+      const [i, v]: number[] = args as number[];
+
+      if (this._dirty === 'back') {
+        await this.read();
+      }
+
+      this.data[i] = v;
+      this._dirty = 'front';
+    } else if (Array.isArray(args[0])) {
+      const [data, oi = 0] = args as [T[], number | undefined];
+
+      if (this._dirty === 'back' && (oi || data.length < this.dims)) {
+        await this.read();
+      }
+
+      let i = oi;
+      for (const v of data ?? []) {
+        if (i >= this.dims) {
+          break;
+        }
+        this.data[i] = v;
+        i++;
+      }
+      this._dirty = 'front';
+    } else {
+      const [vector] = args as [Vector<D, A, T>];
+      const copy = this._device.createCommandEncoder();
+      copy.copyBufferToBuffer(vector.buffer(), 0, this.buffer(), 0, this.byteLength);
+      this._device.queue.submit([copy.finish()]);
+      this._dirty = 'back';
+    }
+  }
+
+  public async get(i: number): Promise<T>;
+  public async get(): Promise<CPUVector<D, T>>;
+  public async get(i?: number): Promise<T | CPUVector<D, T>> {
+    if (this._dirty === 'back') {
+      await this.read();
+    }
+    if (typeof i === 'number') {
+      return this.data[i] as T;
+    } else {
+      const result: CPUVector<D, T> = [] as unknown as CPUVector<D, T>;
+      for (let i: number = 0; i < this.dims; i++) {
+        result.push(this.data[i] as T);
+      }
+      return result;
+    }
+  }
+
+  protected static override assertPackable(pm: number, ...ms: Vector<number, Uint32Array | Int32Array | Float32Array>[]): void {
+    for (const m of ms) {
+      if (m.dims % (Vector.dunit * pm) !== 0) {
+        throw new Error(`Vector size ${m.dims} must be divisible by ${Vector.dunit * pm}.`);
+      }
+    }
+  }
+
+  public static sum<
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Vector<D, A, T>, r: Vector<D, A, T>, result: Vector<D, A, T> = new Vector(l._device, l._array, l.dims)): Vector<D, A, T> {
+    return GPUStruct.csum(l, r, result);
+  }
+
+  public static sub<
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Vector<D, A, T>, r: Vector<D, A, T>, result: Vector<D, A, T> = new Vector(l._device, l._array, l.dims)): Vector<D, A, T> {
+    return GPUStruct.csub(l, r, result);
+  }
+
+  public static mul<
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Scalar<A, T>, r: Vector<D, A, T>, result: Vector<D, A, T> = new Vector(r._device, r._array, r.dims)): Vector<D, A, T> {
+    return GPUStruct.cscale(l, r, result);
+  }
+
+  public static dot<
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Vector<D, A, T>, r: Vector<D, A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> {
+    throw new Error('Not implemented.');
+
+    this.assertUnique(l, r);
+    this.assertUnique(l, result);
+    this.assertColocated(l, r, result);
+    // this.assertPackable(4, r, result);
+
+    return result.calculate('dot', SHADER_DOT, [l, r]);
+  }
+
+  public static unavg<
+    L extends number,
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: L,
+    m: Vector<D, A, T>,
+    result: VectorBatch<L, D, A, T> = new VectorBatch(m._device, m._array, l, m.dims),
+  ): VectorBatch<L, D, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m);
+
+    return result.calculate('unavg', SHADER_UNAVG, [m], false, 4);
+  }
+}

src/data/group-queue.ts

@@ -0,0 +1,100 @@
+import { List, ListItem } from './list';
+
+interface InternalQueueItem<T> {
+  value: T;
+  self?: ListItem<InternalQueueItem<T>>;
+  orphan?: boolean;
+}
+
+export interface QueueItem<T> {
+  readonly value: T;
+  readonly self: ListItem<QueueItem<T>>;
+  readonly orphan?: boolean;
+}
+
+export abstract class GroupQueueBase<K, T> {
+  protected _groups: Map<K, List<InternalQueueItem<T>>> = new Map();
+
+  protected abstract hash(v: T): K;
+
+  public get length(): number {
+    return this._groups.size;
+  }
+
+  public constructor(values?: Iterable<T>) {
+    for (const v of values ?? []) {
+      this.insert(v);
+    }
+  }
+
+  public insert(value: T): QueueItem<T> {
+    const hash: K = this.hash(value);
+    const group: List<InternalQueueItem<T>> = this._groups.get(hash) ?? this._groups.set(hash, new List()).get(hash)!;
+    const item: InternalQueueItem<T> = { value };
+    return (item.self = group.append(item)).value as QueueItem<T>;
+  }
+
+  public remove(item: QueueItem<T>): boolean {
+    if (item.orphan) {
+      return false;
+    }
+    const hash: K = this.hash(item.value);
+    const group: List<InternalQueueItem<T>> | undefined = this._groups.get(hash);
+    if (!group) {
+      return false;
+    }
+    group.delete(item.self);
+    if (!group.length) {
+      this._groups.delete(hash);
+    }
+    return true;
+  }
+
+  public itemOf(value: T): QueueItem<T> | undefined {
+    const hash: K = this.hash(value);
+    const group: List<InternalQueueItem<T>> | undefined = this._groups.get(hash);
+    if (!group) {
+      return undefined;
+    }
+    for (const item of group) {
+      if (item.value === value) {
+        return item as QueueItem<T>;
+      }
+    }
+    return undefined;
+  }
+
+  public peek(): List<QueueItem<T>> | undefined {
+    if (!this._groups.size) {
+      return undefined;
+    }
+
+    let group: List<InternalQueueItem<T>> | undefined = undefined;
+    for (const g of this._groups.values()) {
+      if (g.length > (group?.length ?? 0)) {
+        group = g;
+      }
+    }
+    return group as List<QueueItem<T>>;
+  }
+
+  public pop(): List<QueueItem<T>> | undefined {
+    const result: List<QueueItem<T>> | undefined = this.peek();
+    for (const item of result ?? []) {
+      this.remove(item);
+    }
+    return result;
+  }
+
+  public *[Symbol.iterator](): IterableIterator<List<QueueItem<T>>> {
+    for (const group of [...this._groups.values()].sort(({ length: ll }, { length: rl }) => rl - ll)) {
+      yield group as List<QueueItem<T>>;
+    }
+  }
+}
+
+export class GroupQueue<T> extends GroupQueueBase<T, T> {
+  protected override hash(v: T): T {
+    return v;
+  }
+}

src/data/list.ts

@@ -0,0 +1,127 @@
+interface InternalListItem<T> {
+  value: T;
+  prev: InternalListItem<T>;
+  next: InternalListItem<T>;
+  orphan?: boolean;
+}
+
+export interface ListItem<T> {
+  readonly value: T;
+  readonly prev: ListItem<T>;
+  readonly next: ListItem<T>;
+  readonly orphan?: boolean;
+}
+
+export class List<T> implements Iterable<T> {
+  protected _head?: InternalListItem<T>;
+  protected _length: number = 0;
+
+  public get first(): T | undefined {
+    return this._head?.value;
+  }
+
+  public get last(): T | undefined {
+    return this._head?.prev?.value;
+  }
+
+  public get length(): number {
+    return this._length;
+  }
+
+  public constructor(values?: Iterable<T>) {
+    for (const v of values ?? []) {
+      this.append(v);
+    }
+  }
+
+  public prepend(value: T): ListItem<T> {
+    if (!this._head) {
+      this._head = { value } as InternalListItem<T>;
+      this._head.prev = this._head.next = this._head;
+    } else {
+      this._head = {
+        value,
+        next: this._head,
+        prev: this._head.prev,
+      };
+      this._head.prev.next = this._head.next.prev = this._head;
+    }
+    this._length++;
+    return this._head;
+  }
+
+  public append(value: T): ListItem<T> {
+    this.prepend(value);
+    this._head = this._head?.next;
+    return this._head!.prev;
+  }
+
+  public includes(value: T): boolean {
+    return !!this.itemOf(value);
+  }
+
+  public itemOf(value: T): ListItem<T> | undefined {
+    for (const item of this.items()) {
+      if (item.value === value) {
+        return item;
+      }
+    }
+    return undefined;
+  }
+
+  public set(item: ListItem<T>, value: T): boolean {
+    if (item.orphan) {
+      return false;
+    }
+
+    (item as InternalListItem<T>).value = value;
+
+    return true;
+  }
+
+  public delete(item: ListItem<T>): boolean {
+    if (item.orphan) {
+      return false;
+    }
+
+    if (this._head === item) {
+      this._head = item.next;
+    }
+
+    if (this._head?.next === this._head) {
+      this._head = undefined;
+    }
+
+    (item as InternalListItem<T>).prev.next = item.next;
+    (item as InternalListItem<T>).next.prev = item.prev;
+    (item as InternalListItem<T>).orphan = true;
+
+    this._length--;
+
+    return true;
+  }
+
+  public *[Symbol.iterator](): IterableIterator<T> {
+    for (const item of this.items()) {
+      yield item.value;
+    }
+  }
+
+  public values(): IterableIterator<T> {
+    return this[Symbol.iterator]();
+  }
+
+  public *items(): IterableIterator<ListItem<T>> {
+    if (!this._head) {
+      return;
+    }
+
+    let head = this._head;
+    let item = this._head;
+    do {
+      head = this._head;
+      yield item;
+      item = item.next;
+    } while (item !== head);
+  }
+}

src/data/math.ts

@@ -0,0 +1,142 @@
+export type Vector<D extends number, I = number> = I[] & { length: D };
+export type Matrix<R extends number, C extends number, I = number> = Vector<R, Vector<C, I>>;
+
+export function vec<D extends number, I = number>(s: D, f: (i: number) => I): Vector<D, I> {
+  return Array.from({ length: s }, (_, i) => f(i)) as Vector<D, I>;
+}
+
+export function mat<R extends number, C extends number, I = number>(r: R, c: C, f: (x: number, y: number) => number): Matrix<R, C, I> {
+  return Array.from({ length: r }, (_, y) => Array.from({ length: c }, (_, x) => f(x, y))) as Matrix<R, C, I>;
+}
+
+export function sum<R extends number, C extends number>(l: Matrix<R, C>, r: Matrix<R, C>): Matrix<R, C> {
+  const result: Matrix<R, C> = mat(l.length, l[0].length, () => 0);
+  for (let i = 0; i < l.length; i++) {
+    for (let j = 0; j < r[0].length; j++) {
+      result[i][j] += l[i][j] + r[i][j];
+    }
+  }
+  return result;
+}
+
+export function mul<R extends number, C extends number>(l: Matrix<R, C>, r: number): Matrix<R, C>;
+export function mul<D0 extends number, D1 extends number>(l: Vector<D0>, r: Matrix<D0, D1>): Vector<D1>;
+export function mul<R extends number, S extends number, C extends number>(l: Matrix<R, S>, r: Matrix<S, C>): Matrix<R, C>;
+export function mul(l: Matrix<number, number> | Vector<number>, r: Matrix<number, number> | number): Matrix<number, number> | Vector<number> {
+  if (typeof r === 'number') {
+    const lm: Matrix<number, number> = l as Matrix<number, number>;
+    const result = Array.from({ length: lm.length }, () => Array(lm[0].length).fill(0));
+    for (let i = 0; i < lm.length; i++) {
+      for (let j = 0; j < lm[i].length; j++) {
+        result[i][j] += lm[i][j] * r;
+      }
+    }
+    return result;
+  } else if (!Array.isArray(l[0])) {
+    const lv: Vector<number> = l as Vector<number>;
+    const result = Array(r[0].length).fill(0);
+    for (let j = 0; j < r[0].length; j++) {
+      for (let i = 0; i < lv.length; i++) {
+        result[j] += lv[i] * r[i][j];
+      }
+    }
+
+    return result;
+  } else {
+    const lm: Matrix<number, number> = l as Matrix<number, number>;
+    const result = Array.from({ length: lm.length }, () => Array(r[0].length).fill(0));
+    for (let i = 0; i < lm.length; i++) {
+      for (let j = 0; j < r[0].length; j++) {
+        for (let k = 0; k < r.length; k++) {
+          result[i][j] += lm[i][k] * r[k][j];
+        }
+      }
+    }
+    return result;
+  }
+}
+
+export function transpose<R extends number, C extends number>(m: Matrix<R, C>): Matrix<C, R> {
+  return mat(m[0].length, m.length, (x, y) => m[x][y]);
+}
+
+export function softmax(row: number[]): number[] {
+  const max = Math.max(...row);
+  const exps = row.map(v => Math.exp(v - max));
+  const sum = exps.reduce((a, b) => a + b, 0);
+  return exps.map(v => v / sum);
+}
+
+export function msoftmax<R extends number, C extends number>(m: Matrix<R, C>): Matrix<R, C> {
+  return m.map(softmax) as Matrix<R, C>;
+}
+
+export function norm<R extends number, C extends number>(x: Matrix<R, C>, epsilon = 1e-5): Matrix<R, C> {
+  return x.map(vec => {
+    const mean = vec.reduce((sum, val) => sum + val, 0) / vec.length;
+    const variance = vec.reduce((sum, val) => sum + (val - mean) ** 2, 0) / vec.length;
+    return vec.map(val => (val - mean) / Math.sqrt(variance + epsilon));
+  }) as Matrix<R, C>;
+}
+
+export function relu<D extends number>(vec: Vector<D>): Vector<D> {
+  return vec.map(v => Math.max(0, v)) as Vector<D>;
+}
+
+export function mean<D extends number>(v: Vector<D>): number {
+  return v.reduce((s, c) => s + c, 0) / v.length;
+}
+
+export function drm(data: number[], interval = 0.9): number {
+  const sorted = [...data].sort((a, b) => a - b);
+  const count = Math.floor(data.length * interval);
+  let mi = 0;
+  let mv = Infinity;
+
+  for (let i = 0; i <= data.length - count; i++) {
+    const range = sorted[i + count - 1] - sorted[i];
+    if (range < mv) {
+      mv = range;
+      mi = i;
+    }
+  }
+
+  return sorted.slice(mi, mi + count).reduce((a, b) => a + b, 0) / count;
+}
+
+export function randseq(n: number): number[] {
+  const is = Array.from({ length: n }, (_, i) => i);
+  for (let i = n - 1; i > 0; i--) {
+    const j = Math.floor(Math.random() * (i + 1));
+    [is[i], is[j]] = [is[j], is[i]]; // swap
+  }
+  return is;
+}
+
+export function stddev(xs: number[]): number {
+  const m = mean(xs);
+  return Math.sqrt(xs.reduce((s, v) => s + (v - m) ** 2, 0) / xs.length);
+}
+
+export function ema(xs: number[], alpha = 0.1): number {
+  let result = xs[0];
+  for (let i = 1; i < xs.length; i++) {
+    result = alpha * xs[i] + (1 - alpha) * result;
+  }
+  return result;
+}
+
+export function slope(xs: number[]): number {
+  const xMean = (xs.length - 1) / 2;
+  const yMean = mean(xs);
+
+  let num = 0, den = 0;
+  for (let i = 0; i < xs.length; i++) {
+    const x = i - xMean;
+    const y = xs[i] - yMean;
+    num += x * y;
+    den += x * x;
+  }
+
+  return num / den;
+}

src/data/trie.ts

@@ -0,0 +1,73 @@
+export interface InternalTrieNode {
+  children: Map<string, InternalTrieNode>;
+  token?: string;
+}
+
+export interface TrieNode {
+  readonly children: Map<string, TrieNode>;
+  readonly token?: string;
+}
+
+export class Trie implements Iterable<string> {
+  protected root: InternalTrieNode = { children: new Map() };
+
+  constructor(values?: string[]) {
+    for (const value of values ?? []) {
+      this.insert(value);
+    }
+  }
+
+  public insert(value: string) {
+    let node = this.root;
+    for (const char of value) {
+      node = node.children.get(char) ?? node.children.set(char, { children: new Map() }).get(char)!;
+    }
+    node.token = value;
+  }
+
+  public match(value: string): string | undefined {
+    let node: TrieNode = this.root;
+    let match: string | undefined = undefined;
+    for (let i = 0; i < value.length && node.children.has(value[i]); i++) {
+      match = (node = node.children.get(value[i])!).token ?? match;
+    }
+    return match;
+  }
+
+  public *[Symbol.iterator](): IterableIterator<string> {
+    const queue: InternalTrieNode[] = [];
+    let item: InternalTrieNode | undefined = this.root;
+    while (item) {
+      if (item.token) {
+        yield item.token;
+      }
+      queue.push(...item.children.values());
+      item = queue.pop();
+    }
+  }
+
+  public toJSON(node: InternalTrieNode = this.root): object {
+    return [
+      ...node.token ? [node.token] : [],
+      ...[...node.children.entries()].map(([k, c]) => [k, this.toJSON(c)]),
+    ];
+  }
+
+  protected static nodeFromJSON(element: object): InternalTrieNode {
+    if (!Array.isArray(element) || !element.every(e => typeof e === 'string'
+      || Array.isArray(e) && e.length === 2)) {
+      throw new Error('Malformed Trie JSON.');
+    }
+
+    return {
+      children: new Map((typeof element[0] === 'string' ? element.slice(1) : element).map(([k, v]) => [k, this.nodeFromJSON(v)])),
+      token: typeof element[0] === 'string' ? element[0] : undefined,
+    };
+  }
+
+  public static fromJSON(json: object): Trie {
+    const result = new Trie();
+    result.root = Trie.nodeFromJSON(json);
+    return result;
+  }
+}

src/index.ts

@@ -0,0 +1,196 @@
+import fs from 'node:fs';
+import { create, globals } from 'webgpu';
+import yargs from 'yargs';
+import { Matrix } from './data/gpu/matrix';
+import { Scalar } from './data/gpu/scalar';
+import { mul, sum, transpose } from './data/math';
+import { Model } from './model-gpu';
+import { Tokenizer } from './tokenizer';
+
+Object.assign(globalThis, globals);
+
+const GPU = create([]);
+
+const device = await GPU.requestAdapter({ powerPreference: 'high-performance' }).then(adapter => adapter?.requestDevice());
+
+if (!device) {
+  process.stderr.write('Can not find device.\n');
+  process.exit(1);
+}
+process.stderr.write(`Running on: ${device.adapterInfo.description}.\n`);
+
+const args = await yargs()
+  .option('model', { type: 'string', demandOption: true })
+  .command('test-matrix', 'Test core components')
+  .command('train-tokenizer', 'Train new tokenizer', yargs => yargs
+    .option('vocab', { type: 'number', demandOption: true }))
+  .command('tokenize', 'Tokenize input using existing tokenizer', yargs => yargs
+    .option('input', { type: 'string' })
+    .option('output', { type: 'string' }))
+  .command('configure', 'Configure new model', yargs => yargs
+    .option('window', { type: 'number', demandOption: true })
+    .option('dimensions', { type: 'number', demandOption: true })
+    .option('heads', { type: 'number', demandOption: true })
+    .option('ffs', { type: 'number', demandOption: true }))
+  .command('train', 'Train existing model', yargs => yargs
+    .option('epochs', { type: 'number' })
+    .option('stats', { type: 'string' })
+    .option('logits', { type: 'string' })
+    .option('samples', { type: 'string' }))
+  .command('transform', 'Run single step of existing model', yargs => yargs
+    .option('input', { type: 'string' })
+    .option('output', { type: 'string' })
+    .option('temperature', { type: 'number' })
+    .option('logits', { type: 'string' }))
+  .command('run', 'Run existing model', yargs => yargs
+    .option('input', { type: 'string' })
+    .option('output', { type: 'string' })
+    .option('temperature', { type: 'number' })
+    .option('samples', { type: 'string' })
+    .option('logits', { type: 'string' }))
+  .command('benchmark', 'Benchmark existing model')
+  .demandCommand(1)
+  .parse(process.argv.slice(2));
+
+if (args._.includes('test-matrix')) {
+  (Matrix as any).unit = 1;
+
+  const SL = new Matrix(device, Float32Array, 2, 4, [
+    [3, -7, 6, 12],
+    [11, -5, 7, 3],
+  ]);
+  const SR = new Matrix(device, Float32Array, 2, 4, [
+    [6, 12, 3, -7],
+    [7, 3, 11, -5],
+  ]);
+
+  const SO = new Matrix(device, Float32Array, 2, 4);
+
+  const SS = new Scalar(device, Float32Array, 3);
+
+  const ML = new Matrix(device, Float32Array, 2, 6, [
+    [3, -7, 6, 12, 4, 2],
+    [11, -5, 7, 3, 2, 4],
+  ]);
+  const MR = new Matrix(device, Float32Array, 6, 2, [
+    [6, 12],
+    [3, -7],
+    [7, 3],
+    [11, -5],
+    [4, 2],
+    [2, 4],
+  ]);
+
+  const MO = new Matrix(device, Float32Array, 2, 2);
+
+  console.log('Sum: CPU: ', sum(await SL.get(), await SR.get()));
+  console.log('Sum: GPU (new): ', await Matrix.sum(SL, SR, SO).get());
+  console.log('Sum: GPU (in-place): ', await Matrix.sum(SL, SR, SR).get());
+
+  console.log('Scale: CPU: ', mul(await SL.get(), 3));
+  console.log('Scale: GPU (new): ', await Matrix.scale(SS, SL, SR).get());
+  console.log('Scale: GPU (in-place): ', await Matrix.scale(SS, SL, SL).get());
+
+  console.log('Mul: CPU: ', mul(await ML.get(), await MR.get()));
+  console.log('Mul: GPU: ', await Matrix.mul(ML, MR, MO).get());
+
+  console.log('Transpose: CPU: ', transpose(await ML.get()));
+  console.log('Transpose: GPU: ', await Matrix.transpose(ML, MR).get());
+}
+
+const dir: string = args.model;
+
+let tokenizer: Tokenizer | undefined;
+let model: Model | undefined;
+if (args._.includes('train-tokenizer')) {
+  tokenizer = new Tokenizer(args['vocab'] as number, fs.readFileSync(`${dir}/corpus.txt`).toString());
+
+  fs.writeFileSync(`${dir}/vocab.json`, JSON.stringify(tokenizer.toJSON()));
+  fs.writeFileSync(`${dir}/vocab-flat.txt`, tokenizer.tokens.map(s => JSON.stringify(s).slice(1, -1)).join('\n'));
+}
+
+if (args._.includes('tokenize') || args._.includes('configure') || args._.includes('train') || args._.includes('transform') || args._.includes('run') || args._.includes('benchmark')) {
+  if (!tokenizer && !fs.existsSync(`${dir}/vocab.json`)) {
+    process.stderr.write('Can not find tokenizer data.\n');
+    process.exit(1);
+  }
+
+  tokenizer ??= new Tokenizer(JSON.parse(fs.readFileSync(`${dir}/vocab.json`).toString()));
+
+  if (args._.includes('tokenize')) {
+    fs.writeFileSync(args['output'] as string || 1, tokenizer.detokenize(tokenizer.tokenize(fs.readFileSync(args['input'] as string || 0).toString())).map(s => JSON.stringify(s).slice(1, -1)).join('\n'));
+  }
+
+  if (args._.includes('configure')) {
+    model = new Model(device, tokenizer, {
+      tokenizer: `${tokenizer.tokens.length}${tokenizer.tokens[0]}`,
+      vs: tokenizer.tokens.length,
+      ws: args['window'] as number,
+      ds: args['dimensions'] as number,
+    });
+    fs.writeFileSync(`${dir}/model.json`, JSON.stringify(await model.toJSON()));
+
+    if (args['stats']) { fs.writeFileSync(args['stats'] as string, 'index,timestamp,epoch,batch,lr,loss,accuracy,grad_inp,grad_qry,grad_key,grad_val,grad_out,queueing,computing,inferring\n'); }
+    if (args['logits']) { fs.writeFileSync(args['logits'] as string, 'sample,token,itid,ttid,t0,o0,p0,t1,o1,p1,t2,o2,p2,tNm2,oNm2,pNm2,tNm1,oNm1,pNm1,tN,oN,pN,sum,stddev,entropy,loss\n'); }
+    if (args['samples']) { fs.writeFileSync(args['samples'] as string, ''); }
+  }
+
+  if (args._.includes('train') || args._.includes('transform') || args._.includes('run') || args._.includes('benchmark')) {
+    if (!model && !fs.existsSync(`${dir}/model.json`)) {
+      process.stderr.write('Can not find model data.\n');
+      process.exit(1);
+    }
+
+    model ??= new Model(device, tokenizer, JSON.parse(fs.readFileSync(`${dir}/model.json`).toString()));
+
+    if (args._.includes('train')) {
+      const snapshot = (json: object) => fs.writeFileSync(`${dir}/model.json`, JSON.stringify(json));
+
+      await model.train(
+        args.epochs as number || 'auto',
+        fs.readFileSync(`${dir}/corpus.txt`).toString(),
+        {
+          logs: process.stderr,
+          stats: args['stats'] ? fs.createWriteStream(args['stats'] as string, { flags: 'a' }) : undefined,
+          logits: args['logits'] ? fs.createWriteStream(args['logits'] as string, { flags: 'a' }) : undefined,
+          samples: args['samples'] ? fs.createWriteStream(args['samples'] as string, { flags: 'a' }) : undefined,
+          snapshot,
+        });
+
+      snapshot(await model.toJSON());
+    }
+
+    if (args._.includes('transform')) {
+      const [inp, out] = await model.transform(
+        tokenizer.tokenize(fs.readFileSync(args['input'] as string || 0).toString()),
+        {
+          temp: args['temperature'] as number || 1,
+          logits: args['logits'] ? fs.createWriteStream(args['logits'] as string) : undefined,
+        },
+      );
+
+      fs.writeFileSync(args['output'] as string || 1, `Input:\n${tokenizer.detokenize(inp).join('')}\nOutput:${tokenizer.detokenize(out).join('\n')}\n`);
+    }
+
+    if (args._.includes('run')) {
+      fs.writeFileSync(args['output'] as string || 1, tokenizer.detokenize(await model.run(
+        tokenizer.tokenize(fs.readFileSync(args['input'] as string || 0).toString()),
+        {
+          temp: args['temperature'] as number || 1,
+          logits: args['logits'] ? fs.createWriteStream(args['logits'] as string) : undefined,
+          samples: args['samples'] as number || model.ws,
+        },
+      )).join('') + '\n');
+    }
+
+    if (args._.includes('benchmark')) {
+      throw new Error('Not implemented.');
+    }
+  }
+}
+
+await device.queue.onSubmittedWorkDone();
+
+device.destroy();
+
+process.stderr.write('All done!\n');

src/model-cpu.ts

@@ -0,0 +1,261 @@
+import { Matrix, Vector, mat, mul, softmax, sum, transpose } from './data/math';
+import { TID } from './tokenizer';
+
+export interface ModelInitData<
+  V extends number,
+  W extends number,
+  D extends number,
+  H extends number,
+  F extends number,
+> {
+  readonly vs: V;
+  readonly ws: W;
+  readonly ds: D;
+  readonly hs: H;
+  readonly fs: F;
+}
+
+export interface ModelData<
+  V extends number,
+  W extends number,
+  D extends number,
+  H extends number,
+  HD extends number,
+  F extends number,
+  FD extends number,
+> extends ModelInitData<V, W, D, H, F> {
+  readonly trained: true;
+  readonly positions: Matrix<W, D>;
+  readonly embeddings: Matrix<V, D>;
+  readonly weights: {
+    query: Matrix<D, HD>;
+    key: Matrix<D, HD>;
+    value: Matrix<D, HD>;
+    attn: Matrix<D, D>;
+    ffi: Matrix<D, FD>;
+    ffo: Matrix<FD, D>;
+    out: Matrix<D, V>;
+  };
+}
+
+const ES: number = 1;
+
+export class Model<
+  V extends number = number,
+  W extends number = number,
+  D extends number = number,
+  H extends number = number,
+  HD extends number = number,
+  F extends number = number,
+  FD extends number = number,
+> {
+  protected readonly vs: V;
+  protected readonly ws: W;
+  protected readonly ds: D;
+  protected readonly hs: H;
+  protected readonly hds: HD;
+  protected readonly fs: F;
+  protected readonly fds: FD;
+
+  protected positions: Matrix<W, D>;
+  protected embeddings: Matrix<V, D>;
+
+  protected readonly weights: {
+    query: Matrix<D, HD>;
+    key: Matrix<D, HD>;
+    value: Matrix<D, HD>;
+    attn: Matrix<D, D>;
+    ffi: Matrix<D, FD>;
+    ffo: Matrix<FD, D>;
+    out: Matrix<D, V>;
+  };
+
+  protected embed(tids: Vector<W, TID>): Matrix<W, D> {
+    return sum(tids.map(tid => this.embeddings[tid]) as Matrix<W, D>, this.positions);
+  }
+
+  // protected attentize<N extends number>(input: Matrix<N, D>): Matrix<N, D> {
+  //   const SCALE: number = 1 / Math.sqrt(this.hds);
+  //   const Q: Matrix<N, HD> = mul(input, this.weights.query);
+  //   const K: Matrix<N, HD> = mul(input, this.weights.key);
+  //   const V: Matrix<N, HD> = mul(input, this.weights.value);
+  //   const scores: Matrix<N, N> = msoftmax(mul(mul(Q, transpose(K)), SCALE));
+  //   const out: Matrix<N, HD> = mul(scores, V);
+  //   return mul(out as unknown as Matrix<N, D>, this.weights.attn);
+  // }
+
+  // protected feedforward<N extends number>(input: Matrix<N, D>): Matrix<N, D> {
+  //   return input.map((v: Vector<D>) => mul(relu(mul(v, this.weights.ffi)), this.weights.ffo)) as Matrix<N, D>;
+  // }
+
+  // protected transform<N extends number>(input: Matrix<N, D>): Matrix<N, D> {
+  //   const attned: Matrix<N, D> = sum(input, this.attentize(input));
+  //   const ffned: Matrix<N, D> = sum(input, this.feedforward(attned));
+  //   return ffned;
+  // }
+
+  protected logitize(input: Matrix<W, D>): Matrix<W, V> {
+    return mul(input, this.weights.out).map(softmax) as Matrix<W, V>;
+  }
+
+  protected argmax(logit: Vector<V>): TID {
+    return logit.reduce((mi, p, i, a) => p > a[mi] ? i : mi, 0);
+  }
+
+  protected sample(logit: Vector<V>): TID {
+    let p: number = 0;
+    const t: number = Math.random();
+    for (let i = 0; i < logit.length; i++) {
+      p += logit[i];
+      if (t < p) {
+        return i;
+      }
+    }
+    return 0;
+  }
+
+  protected forward(tids: Vector<W, TID>): Matrix<W, V> {
+    return this.logitize(this.embed(tids));
+  }
+
+  protected loss<N extends number>(probs: Vector<N, Vector<V>>, target: Vector<N, TID>): number {
+    return probs.reduce((l: number, ps: Vector<V>, i: number) => l - Math.log(ps[target[i]] + 1e-10), 0) / probs.length;
+  }
+
+  protected grad(probs: Matrix<W, V>, target: Vector<W, TID>): Matrix<W, V> {
+    const grads = mul(probs, 1);
+    target.forEach((t, i) => grads[i][t] -= 1);
+    return grads;
+  }
+
+  protected updateWOut<N extends number>(input: Matrix<N, D>, grad: Matrix<N, V>, lr: number): void {
+    this.weights.out = sum(mul(mul(transpose(input), grad), -lr), this.weights.out);
+  }
+
+  public constructor(_device: unknown, data: ModelInitData<V, W, D, H, F> | ModelData<V, W, D, H, HD, F, FD>) {
+    this.vs = data.vs;
+    this.ws = data.ws;
+    this.ds = data.ds;
+    this.hs = data.hs;
+    this.hds = (this.ds / this.hs) as HD;
+    this.fs = data.fs;
+    this.fds = (this.ds * this.fs) as FD;
+    if (!('trained' in data)) {
+      this.positions = mat(this.ws, this.ds, (d, i) => {
+        const angle = i / Math.pow(10000, 2 * Math.floor(d / 2) / this.ds);
+        return d % 2 === 0 ? Math.sin(angle) * ES : Math.cos(angle) * ES;
+      });
+      this.embeddings = mat(this.vs, this.ds, () => Math.random() * 2 * ES - ES);
+      this.weights = {
+        query: mat(this.ds, this.hds, () => Math.random() * 0.02 - 0.01),
+        key: mat(this.ds, this.hds, () => Math.random() * 0.02 - 0.01),
+        value: mat(this.ds, this.hds, () => Math.random() * 0.02 - 0.01),
+        attn: mat(this.ds, this.ds, () => Math.random() * 0.02 - 0.01),
+        ffi: mat(this.ds, this.fds, () => Math.random() * 0.02 - 0.01),
+        ffo: mat(this.fds, this.ds, () => Math.random() * 0.02 - 0.01),
+        out: mat(this.ds, this.vs, () => Math.random() * 0.2 - 0.1),
+      };
+    } else {
+      this.positions = data.positions;
+      this.embeddings = data.embeddings;
+      this.weights = data.weights;
+    }
+  }
+
+  public train(data: TID[], epochs: number): void {
+    let log = Date.now();
+    const samples: number[] = [];
+    console.log('Training model.');
+
+    const runs = data.length - this.ws - 1;
+    for (let epoch = 0; epoch < epochs; epoch++) {
+      for (let s = 0; s < runs; s++) {
+        const begin: number = Date.now();
+
+        const source = data.slice(s, s + this.ws) as Vector<W>;
+        const target = data.slice(s + 1, s + 1 + this.ws) as Vector<W>;
+
+        // const encoded = transform(model, embed(model, source));
+
+        const embedded = this.embed(source);
+        const logits = this.logitize(embedded);
+
+        this.updateWOut(embedded, this.grad(logits, target), 0.001);
+
+        const output = logits.map(this.sample);
+
+        const a = 100 * target.filter((t, i) => output[i] === t).length / target.length;
+        const l = this.loss(logits, target);
+
+        const end: number = Date.now();
+        samples.push(end - begin);
+        if ((epoch * runs + s) % Math.round(runs * epochs / Math.min(100, epochs)) === 0 || epoch === epochs - 1 && s === runs - 1 || end - log > 5000) {
+          console.log(`Done ${samples.length} / ${epochs * runs}, accuracy ${a.toFixed(0)}%, loss ${l}, in ${end - begin}ms, ETA ${Math.round(samples.reduce((a, b) => a + b, 0) / samples.length * (epochs * runs - samples.length) / 1000)}s`);
+          log = end;
+        }
+      }
+    }
+  }
+
+  public run(tids: TID[], samples: number = this.ws): TID[] {
+    const inputs: Vector<W> = [
+      ...Array(Math.max(0, this.ws - tids.length)).fill(0),
+      ...tids.slice(-this.ws)] as Vector<W>;
+    const result: Vector<W, TID> = [] as unknown as Vector<W, TID>;
+    while (result.length < samples) {
+      const next: TID = this.sample(this.forward(inputs)[this.ws - 1]);
+      result.push(next);
+      inputs.shift();
+      inputs.push(next);
+    }
+    return result;
+  }
+
+  public benchmark(data: TID[]): void {
+    let log = Date.now();
+    const samples: number[] = [];
+    console.log('Testing model.');
+
+    let accuracy: number = 0;
+    let loss: number = 0;
+
+    const runs = data.length - this.ws - 1;
+    for (let s = 0; s < runs; s++) {
+      const begin: number = Date.now();
+
+      const source = data.slice(s, s + this.ws) as Vector<W>;
+      const target = data.slice(s + 1, s + 1 + this.ws) as Vector<W>;
+
+      const embedded = this.embed(source);
+      const logits = this.logitize(embedded);
+      const output = logits.map(this.sample);
+
+      const a = 100 * target.filter((t, i) => output[i] === t).length / target.length;
+      const l = this.loss(logits, target);
+
+      accuracy += a / runs;
+      loss += l / runs;
+
+      const end: number = Date.now();
+      samples.push(end - begin);
+      if (end - log > 5000 || s === runs - 1) {
+        console.log(`Done ${samples.length} / ${runs}, accuracy ${a.toFixed(0)}% / ${accuracy.toFixed()}%, loss ${l.toFixed(4)} / ${loss.toFixed(4)}, in ${end - begin}ms, ETA ${Math.round(samples.reduce((a, b) => a + b, 0) / samples.length * (runs - samples.length) / 1000)}s`);
+        log = end;
+      }
+    }
+  }
+
+  public toJSON(): ModelData<V, W, D, H, HD, F, FD> {
+    return {
+      vs: this.vs,
+      ws: this.ws,
+      ds: this.ds,
+      hs: this.hs,
+      fs: this.fs,
+      trained: true,
+      positions: this.positions,
+      embeddings: this.embeddings,
+      weights: this.weights,
+    };
+  }
+}

src/model-gpu.ts

@@ -0,0 +1,662 @@
+import fs from 'node:fs';
+import { Writable } from 'node:stream';
+import { Matrix } from './data/gpu/matrix';
+import { Scalar } from './data/gpu/scalar';
+import { Vector } from './data/gpu/vector';
+import { Matrix as CPUMatrix, Vector as CPUVector, drm, mat, randseq, slope, stddev, vec } from './data/math';
+import { TID, Tokenizer } from './tokenizer';
+
+const SHADER_HOTONES = fs.readFileSync('shader/hotones.wgsl').toString();
+const SHADER_EMBED = fs.readFileSync('shader/embed.wgsl').toString();
+const SHADER_ATTW = fs.readFileSync('shader/attw.wgsl').toString();
+const SHADER_SOFTMAX_EXP = fs.readFileSync('shader/softmaxexp.wgsl').toString();
+const SHADER_SOFTMAX_NORM = fs.readFileSync('shader/softmaxnorm.wgsl').toString();
+const SHADER_DLDO = fs.readFileSync('shader/dldo.wgsl').toString();
+const SHADER_OTDLDO = fs.readFileSync('shader/otdldo.wgsl').toString();
+const SHADER_TOTDLDOWA = fs.readFileSync('shader/totdldowa.wgsl').toString();
+const SHADER_DLDS = fs.readFileSync('shader/dlds.wgsl').toString();
+const SHADER_GRAD_TOUT = fs.readFileSync('shader/gradtout.wgsl').toString();
+const SHADER_GRAD_TATTW = fs.readFileSync('shader/gradtattw.wgsl').toString();
+const SHADER_GRAD_INP = fs.readFileSync('shader/gradinp.wgsl').toString();
+const SHADER_ARGMAX = fs.readFileSync('shader/argmax.wgsl').toString();
+const SHADER_SAMPLES = fs.readFileSync('shader/samples.wgsl').toString();
+
+export interface ModelInitData<
+  V extends number,
+  W extends number,
+  D extends number,
+> {
+  readonly tokenizer: string;
+  readonly vs: V;
+  readonly ws: W;
+  readonly ds: D;
+}
+
+export interface ModelData<
+  V extends number,
+  W extends number,
+  D extends number,
+> extends ModelInitData<V, W, D> {
+  stage: 'early' | 'mid' | 'end';
+  tokenizer: string;
+  samples: number;
+  weights: {
+    inp: CPUMatrix<V, D>;
+    pos: CPUMatrix<W, D>;
+    tqry: CPUMatrix<D, D>;
+    tkey: CPUMatrix<D, D>;
+    tval: CPUMatrix<D, D>;
+    out: CPUMatrix<V, D>;
+  };
+}
+
+interface TrainingPoint {
+  index: number;
+  timestamp: number;
+  epoch: number;
+  batch: number;
+  lr: number;
+  loss: number;
+  ginp: number;
+  gqry: number;
+  gkey: number;
+  gval: number;
+  gout: number;
+  accuracy: number;
+  queueing: number;
+  computing: number;
+  inferring: number;
+}
+
+interface LogitsPoint {
+  sample: number;
+  token: number;
+  itid: string;
+  ttid: string;
+  t0: string;
+  o0: number;
+  p0: number;
+  t1: string;
+  o1: number;
+  p1: number;
+  t2: string;
+  o2: number;
+  p2: number;
+  tNm2: string;
+  oNm2: number;
+  pNm2: number;
+  tNm1: string;
+  oNm1: number;
+  pNm1: number;
+  tN: string;
+  oN: number;
+  pN: number;
+  sum: number;
+  stddev: number;
+  entropy: number;
+  loss: number;
+}
+
+export class Model<
+  V extends number = number,
+  W extends number = number,
+  D extends number = number,
+> {
+  protected readonly device: GPUDevice;
+  public readonly tokenizer: Tokenizer;
+  public readonly vs: V;
+  public readonly ws: W;
+  public readonly ds: D;
+
+  protected samples: number = 0;
+  protected stage: 'early' | 'mid' | 'end' = 'early';
+
+  protected readonly weights: {
+    readonly inp: Matrix<V, D, Float32Array>;
+    readonly pos: Matrix<W, D, Float32Array>;
+    readonly tqry: Matrix<D, D, Float32Array>;
+    readonly tkey: Matrix<D, D, Float32Array>;
+    readonly tval: Matrix<D, D, Float32Array>;
+    readonly tout: Matrix<V, D, Float32Array>;
+  };
+
+  protected readonly processing: {
+    readonly ds: Scalar<Float32Array>;
+    readonly lr: Scalar<Float32Array>;
+    readonly itids: Vector<W, Int32Array, TID>;
+    readonly ilogits: Matrix<W, V, Float32Array>;
+    readonly e: Matrix<W, D, Float32Array>;
+    readonly q: Matrix<W, D, Float32Array>;
+    readonly k: Matrix<W, D, Float32Array>;
+    readonly tk: Matrix<D, W, Float32Array>;
+    readonly qtk: Matrix<W, W, Float32Array>;
+    readonly v: Matrix<W, D, Float32Array>;
+    readonly s: Matrix<W, W, Float32Array>;
+    readonly se: Matrix<W, W, Float32Array>;
+    readonly wa: Matrix<W, W, Float32Array>;
+    readonly a: Matrix<W, D, Float32Array>;
+    readonly o: Matrix<W, V, Float32Array>;
+    readonly temp: Scalar<Float32Array>;
+    readonly tempo: Matrix<W, V, Float32Array>;
+    readonly on: Matrix<W, V, Float32Array>;
+    readonly ologits: Matrix<W, V, Float32Array>;
+    readonly dldo: Matrix<W, V, Float32Array>;
+    readonly otdldo: Matrix<D, W, Float32Array>;
+    readonly totdldowa: Matrix<W, D, Float32Array>;
+    readonly dlds: Matrix<W, W, Float32Array>;
+    readonly dldsk: Matrix<W, D, Float32Array>;
+    readonly dldsq: Matrix<W, D, Float32Array>;
+    readonly ginp: Matrix<V, D, Float32Array>;
+    readonly gtqry: Matrix<D, D, Float32Array>;
+    readonly gtkey: Matrix<D, D, Float32Array>;
+    readonly gtval: Matrix<D, D, Float32Array>;
+    readonly gtout: Matrix<V, D, Float32Array>;
+    readonly rnd: Vector<W, Float32Array>;
+    readonly otids: Vector<W, Int32Array, TID>;
+    readonly ttids: Vector<W, Int32Array, TID>;
+  };
+
+  protected softmax<H extends number, W extends number>(
+    m: Matrix<H, W, Float32Array>,
+    e: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
+    r: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
+  ): Matrix<H, W, Float32Array> {
+    e.calculate('softmaxexp', SHADER_SOFTMAX_EXP, [m], false, 4);
+    r.calculate('softmaxnorm', SHADER_SOFTMAX_NORM, [e], false, 4);
+    return r;
+  }
+
+  protected async ilogits(itids: Vector<W, Int32Array, TID> = this.processing.itids): Promise<Matrix<W, V, Float32Array>> {
+    this.processing.ilogits.calculate('hotones', SHADER_HOTONES, [itids]);
+    return this.processing.ilogits;
+  }
+
+  protected async e(ilogits: Matrix<W, V, Float32Array> = this.processing.ilogits): Promise<Matrix<W, D, Float32Array>> {
+    this.processing.e.calculate('embed', SHADER_EMBED, [ilogits, this.weights.inp, this.weights.pos]);
+    return this.processing.e;
+  }
+
+  protected q(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> {
+    this.processing.q.calculate('attw', SHADER_ATTW, [e, this.weights.tqry]);
+    return this.processing.q;
+  }
+
+  protected k(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> {
+    this.processing.k.calculate('attw', SHADER_ATTW, [e, this.weights.tkey]);
+    return this.processing.k;
+  }
+
+  protected v(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> {
+    this.processing.v.calculate('attw', SHADER_ATTW, [e, this.weights.tval]);
+    return this.processing.v;
+  }
+
+  protected s(
+    q: Matrix<W, D, Float32Array> = this.processing.q,
+    k: Matrix<W, D, Float32Array> = this.processing.k,
+  ): Matrix<W, W, Float32Array> {
+    Matrix.transpose(k, this.processing.tk);
+    Matrix.mul(q, this.processing.tk, this.processing.qtk);
+    Matrix.scale(this.processing.ds, this.processing.qtk, this.processing.s);
+    return this.processing.s;
+  }
+
+  protected async wa(s: Matrix<W, W, Float32Array> = this.processing.s): Promise<Matrix<W, W, Float32Array>> {
+    this.softmax(s, this.processing.se, this.processing.wa);
+    return this.processing.wa;
+  }
+
+  protected a(
+    wa: Matrix<W, W, Float32Array> = this.processing.wa,
+    v: Matrix<W, D, Float32Array> = this.processing.v,
+  ): Matrix<W, D, Float32Array> {
+    Matrix.mul(wa, v, this.processing.a);
+    return this.processing.a;
+  }
+
+  protected o(a: Matrix<W, D, Float32Array> = this.processing.a): Matrix<W, V, Float32Array> {
+    this.processing.o.calculate('attw', SHADER_ATTW, [a, this.weights.tout]);
+    return this.processing.o;
+  }
+
+  protected async ologits(o: Matrix<W, V, Float32Array> = this.processing.o, temp: number = 1): Promise<Matrix<W, V, Float32Array>> {
+    await this.processing.temp.set(temp);
+    Matrix.scale(this.processing.temp, o, this.processing.tempo);
+    this.softmax(this.processing.tempo, this.processing.on, this.processing.ologits);
+    return this.processing.ologits;
+  }
+
+  protected async rnd(): Promise<Vector<W, Float32Array>> {
+    await this.processing.rnd.set(vec(this.ws, () => Math.random()));
+    return this.processing.rnd;
+  }
+
+  protected otids(): Vector<W, Int32Array, TID>;
+  protected otids(ologits: Matrix<W, V, Float32Array>): Vector<W, Int32Array, TID>;
+  protected otids(rnd: Vector<W, Float32Array>, ologits: Matrix<W, V, Float32Array>): Vector<W, Int32Array, TID>;
+  protected otids(...args: unknown[]): Vector<W, Int32Array, TID> {
+    if (args.length === 1) {
+      const [ologits = this.processing.ologits] = args as [Matrix<W, V, Float32Array>];
+      this.processing.otids.calculate('argmax', SHADER_ARGMAX, [ologits]);
+      return this.processing.otids;
+    } else {
+      const [rnd = this.processing.rnd, ologits = this.processing.ologits] = args as [Vector<W, Float32Array> | undefined, Matrix<W, V, Float32Array> | undefined];
+      this.processing.otids.calculate('samples', SHADER_SAMPLES, [ologits, rnd]);
+      return this.processing.otids;
+    }
+  }
+
+  protected loss<N extends number>(probs: CPUVector<N, CPUVector<V>>, target: CPUVector<N, TID>): number {
+    return probs.reduce((l: number, ps: CPUVector<V>, i: number) => l - Math.log(ps[target[i]] + 1e-10), 0) / probs.length;
+  }
+
+  protected async upd(
+    ilogits: Matrix<W, V, Float32Array> = this.processing.ilogits,
+    e: Matrix<W, D, Float32Array> = this.processing.e,
+    q: Matrix<W, D, Float32Array> = this.processing.q,
+    k: Matrix<W, D, Float32Array> = this.processing.k,
+    v: Matrix<W, D, Float32Array> = this.processing.v,
+    wa: Matrix<W, W, Float32Array> = this.processing.wa,
+    a: Matrix<W, D, Float32Array> = this.processing.a,
+    ologits: Matrix<W, V, Float32Array> = this.processing.ologits,
+    ttids: Vector<W, Int32Array, TID> = this.processing.ttids,
+    lr: number = 0.001,
+  ): Promise<void> {
+    await this.processing.lr.set(-lr);
+
+    this.processing.dldo.calculate('dldo', SHADER_DLDO, [
+      ologits,
+      ttids,
+      this.processing.lr,
+    ]);
+
+    this.processing.otdldo.calculate('otdldo', SHADER_OTDLDO, [this.processing.dldo, this.weights.tout]);
+
+    this.processing.dlds.calculate('dlds', SHADER_DLDS, [
+      this.processing.ds,
+      wa,
+      v,
+      this.processing.otdldo,
+    ]);
+
+    Matrix.mul(this.processing.dlds, k, this.processing.dldsk);
+    Matrix.mul(Matrix.transpose(this.processing.dlds), q, this.processing.dldsq);
+
+    this.processing.ginp.calculate('gradinp', SHADER_GRAD_INP, [
+      ilogits,
+      this.processing.dldsq,
+      this.processing.dldsk,
+      this.processing.totdldowa,
+      this.weights.tqry,
+      this.weights.tkey,
+      this.weights.tval,
+    ]);
+    Matrix.sum(this.processing.ginp, this.weights.inp, this.weights.inp);
+
+    this.processing.gtout.calculate('gradtout', SHADER_GRAD_TOUT, [a, this.processing.dldo]);
+    Matrix.sum(this.processing.gtout, this.weights.tout, this.weights.tout);
+    this.processing.gtqry.calculate('gradtqry', SHADER_GRAD_TATTW, [e, this.processing.dldsk]);
+    Matrix.sum(this.processing.gtqry, this.weights.tqry, this.weights.tqry);
+    this.processing.gtkey.calculate('gradtkey', SHADER_GRAD_TATTW, [e, this.processing.dldsq]);
+    Matrix.sum(this.processing.gtkey, this.weights.tkey, this.weights.tkey);
+
+    this.processing.totdldowa.calculate('totdldowa', SHADER_TOTDLDOWA, [wa, this.processing.otdldo]);
+    this.processing.gtval.calculate('gradtval', SHADER_GRAD_TATTW, [e, this.processing.totdldowa]);
+    Matrix.sum(this.processing.gtval, this.weights.tval, this.weights.tval);
+  }
+
+  public constructor(
+    device: GPUDevice,
+    tokenizer: Tokenizer,
+    data: ModelInitData<V, W, D> | ModelData<V, W, D>,
+  ) {
+    this.device = device;
+    this.tokenizer = tokenizer;
+    this.vs = data.vs;
+    this.ws = data.ws;
+    this.ds = data.ds;
+    if (!('stage' in data)) {
+      this.weights = {
+        inp: new Matrix(this.device, Float32Array, mat(this.vs, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))),
+        pos: new Matrix(this.device, Float32Array, mat(this.ws, this.ds, (i, d) => {
+          const angle = i / Math.pow(10000, 2 * Math.floor(d / 2) / this.ds);
+          return (d % 2 === 0 ? Math.sin(angle) : Math.cos(angle)) * Math.sqrt(1 / this.ds);
+        })),
+        tqry: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))),
+        tkey: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))),
+        tval: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))),
+        tout: new Matrix(this.device, Float32Array, mat(this.vs, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))),
+      };
+    } else {
+      this.weights = {
+        inp: new Matrix(this.device, Float32Array, data.weights.inp),
+        pos: new Matrix(this.device, Float32Array, data.weights.pos),
+        tqry: new Matrix(this.device, Float32Array, data.weights.tqry),
+        tkey: new Matrix(this.device, Float32Array, data.weights.tkey),
+        tval: new Matrix(this.device, Float32Array, data.weights.tval),
+        tout: new Matrix(this.device, Float32Array, data.weights.out),
+      };
+      this.stage = data.stage;
+      this.samples = data.samples;
+      if (data.tokenizer !== `${tokenizer.tokens.length}${tokenizer.tokens[0]}`) {
+        throw new Error(`The model was trained against a different tokenizer ("${data.tokenizer}").`);
+      }
+    }
+
+    this.processing = {
+      ds: new Scalar<Float32Array>(this.device, Float32Array, 1 / Math.sqrt(this.ds)),
+      lr: new Scalar<Float32Array>(this.device, Float32Array, 0),
+      itids: new Vector<W, Int32Array, TID>(this.device, Int32Array, this.ws),
+      ilogits: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
+      e: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      q: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      k: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      tk: new Matrix<D, W, Float32Array>(this.device, Float32Array, this.ds, this.ws),
+      qtk: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
+      v: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      s: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
+      se: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
+      wa: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
+      a: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      o: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
+      temp: new Scalar<Float32Array>(this.device, Float32Array, 1),
+      tempo: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
+      on: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
+      ologits: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
+      dldo: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
+      otdldo: new Matrix<D, W, Float32Array>(this.device, Float32Array, this.ds, this.ws),
+      totdldowa: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      dlds: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
+      dldsk: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      dldsq: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      ginp: new Matrix(this.device, Float32Array, this.vs, this.ds),
+      gtqry: new Matrix(this.device, Float32Array, this.ds, this.ds),
+      gtkey: new Matrix(this.device, Float32Array, this.ds, this.ds),
+      gtval: new Matrix(this.device, Float32Array, this.ds, this.ds),
+      gtout: new Matrix(this.device, Float32Array, this.vs, this.ds),
+      rnd: new Vector<W, Float32Array>(this.device, Float32Array, this.ws),
+      otids: new Vector<W, Int32Array, TID>(this.device, Int32Array, this.ws),
+      ttids: new Vector<W, Int32Array, TID>(this.device, Int32Array, this.ws),
+    };
+  }
+
+  public async debug(
+    sample: number,
+    itids: Vector<W, Int32Array, TID> = this.processing.itids,
+    ttids: Vector<W, Int32Array, TID> = this.processing.ttids,
+    o: Matrix<W, V, Float32Array> = this.processing.o,
+    ologits: Matrix<W, V, Float32Array> = this.processing.ologits,
+    file?: Writable,
+    llo: boolean = false,
+  ): Promise<void> {
+    if (!file) {
+      return;
+    }
+
+    for (let token = llo ? itids.dims - 1 : 0; token < itids.dims; token++) {
+      const os = (await o.get())[token];
+      const ls = (await ologits.get())[token]
+        .map((p, tid) => [p, tid])
+        .sort(([l], [r]) => r - l);
+
+      const point: LogitsPoint = {
+        sample,
+        token,
+        itid: this.tokenizer.detokenize([await itids.get(token)])[0],
+        ttid: this.tokenizer.detokenize([await ttids.get(token)])[0],
+        t0: this.tokenizer.detokenize([ls[0][1]])[0],
+        o0: os[ls[0][1]],
+        p0: ls[0][0],
+        t1: this.tokenizer.detokenize([ls[1][1]])[0],
+        o1: os[ls[1][1]],
+        p1: ls[1][0],
+        t2: this.tokenizer.detokenize([ls[2][1]])[0],
+        o2: os[ls[2][1]],
+        p2: ls[2][0],
+        tNm2: this.tokenizer.detokenize([ls[ls.length - 3][1]])[0],
+        oNm2: os[ls[ls.length - 3][1]],
+        pNm2: ls[ls.length - 3][0],
+        tNm1: this.tokenizer.detokenize([ls[ls.length - 2][1]])[0],
+        oNm1: os[ls[ls.length - 3][1]],
+        pNm1: ls[ls.length - 2][0],
+        tN: this.tokenizer.detokenize([ls[ls.length - 1][1]])[0],
+        oN: os[ls[ls.length - 1][1]],
+        pN: ls[ls.length - 1][0],
+        sum: (await ologits.get())[token].reduce((a, b) => a + b, 0),
+        stddev: stddev((await o.get())[token]),
+        entropy: -(await ologits.get())[token].reduce((sum, p) => sum + (p > 0 ? p * Math.log(p) : 0), 0),
+        loss: -Math.log((await ologits.get())[token][await ttids.get(token)] + 1e-10),
+      };
+      file.write(`${JSON.stringify(Object.values(point)).slice(1, -1)}\n`);
+    }
+  }
+
+  public async train(
+    duration: number | 'auto',
+    data: string,
+    { logs = process.stderr, stats, logits, samples, snapshot, statsInterval = 5000, snapshotInterval = 5000 }: {
+      logs?: Writable;
+      stats?: Writable;
+      logits?: Writable;
+      samples?: Writable;
+      snapshot?: (json: object) => void;
+      statsInterval?: number;
+      snapshotInterval?: number;
+    } = {},
+  ): Promise<void> {
+    logs.write('Training model...\n');
+
+    let logged = Date.now();
+    let snapshotted = Date.now();
+
+    const aw: number = 10;
+
+    const times = [];
+    const losses = [];
+
+    let epochs = duration === 'auto' ? Infinity : +duration;
+    const tids = this.tokenizer.tokenize(data);
+    const batches = tids.length - this.ws - 1;
+
+    let lr = 0.003;
+
+    for (let epoch = 0; epoch < epochs; epoch++) {
+      const idxs = randseq(tids.length);
+      for (let batch = 0; batch < batches; batch++) {
+        const begin: number = Date.now();
+
+        const si = idxs[batch];
+        await this.processing.itids.set(tids.slice(si, si + this.ws) as CPUVector<W>);
+        await this.processing.ttids.set(tids.slice(si + 1, si + 1 + this.ws) as CPUVector<W>);
+
+        const ilogits = await this.ilogits(this.processing.itids);
+        const e = await this.e(ilogits);
+        const q = this.q(e);
+        const k = this.k(e);
+        const v = this.v(e);
+        const s = this.s(q, k);
+        const wa = await this.wa(s);
+        const a = this.a(wa, v);
+        const o = this.o(a);
+        const ologits = await this.ologits(o);
+        await this.upd(ilogits, e, q, k, v, wa, a, ologits, this.processing.ttids, lr);
+
+        const queued: number = Date.now();
+
+        await this.device.queue.onSubmittedWorkDone();
+
+        const end: number = Date.now();
+
+        this.samples += 1;
+        times.push(end - begin);
+
+        if (
+          epoch === 0 && batch === 0
+          || batch === batches - 1
+          || end - logged > statsInterval
+        ) {
+          const point: TrainingPoint = {
+            index: this.samples,
+            timestamp: Date.now(),
+            epoch,
+            batch,
+            lr,
+            loss: 0,
+            accuracy: 0,
+            ginp: (await this.processing.ginp.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
+            gqry: (await this.processing.gtqry.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
+            gkey: (await this.processing.gtkey.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
+            gval: (await this.processing.gtval.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
+            gout: (await this.processing.gtout.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
+            queueing: 0,
+            computing: 0,
+            inferring: 0,
+          };
+          point.queueing = Date.now() - begin;
+          point.computing = Date.now() - queued;
+
+          const out = await this.otids(await this.rnd(), ologits).get();
+          const tgt = await this.processing.ttids.get();
+
+          const inferred: number = Date.now();
+
+          point.inferring = inferred - end;
+
+          point.accuracy = tgt.filter((t, i) => out[i] === t).length / tgt.length;
+
+          point.loss = this.loss(await ologits.get(), tgt);
+
+          losses.push(point.loss);
+
+          if (this.stage === 'early') {
+            lr = Math.max(0.001, Math.min(0.003, point.loss * 0.001));
+
+            if (losses.length > aw && (slope(losses.slice(-aw)) > -0.005 || stddev(losses.slice(-aw)) < 0.2)) {
+              logs?.write('Switching to mid stage.\n');
+              this.stage = 'mid';
+            }
+          } else if (this.stage === 'mid') {
+            lr = Math.max(0.0003, lr * 0.98);
+
+            if (losses.length > aw && Math.abs(slope(losses.slice(-aw))) < 0.002 && stddev(losses.slice(-aw)) < 0.05) {
+              logs?.write('Switching to end stage.\n');
+              this.stage = 'end';
+            }
+          } else {
+            lr = Math.max(1e-4, lr * 0.98);
+          }
+
+          logs?.write(`Done ${epoch * batches + batch + 1} / ${epochs * batches}, ETA ${Math.round(drm(times) * (epochs * batches - times.length) / 1000)}s\n`);
+          stats?.write(`${Object.values(point).join(',')}\n`);
+          await this.debug(this.samples, this.processing.itids, this.processing.ttids, o, ologits, logits);
+
+          logged = end;
+        }
+
+        if (batch === 0 && end - snapshotted > snapshotInterval) {
+          if (duration === 'auto' && this.stage === 'end' && losses.length > aw && (
+            Math.abs(slope(losses.slice(-aw))) < 0.001 && stddev(losses.slice(-aw)) < 0.001
+            || slope(losses.slice(-aw)) > 0.001
+            || stddev(losses.slice(-aw)) > 1)) {
+            logs?.write(`Ending training with slope=${slope(losses.slice(-aw))} stddev=${stddev(losses.slice(-aw))}.\n`);
+            epochs = 0;
+            break;
+          }
+
+          snapshot?.(await this.toJSON());
+
+          samples?.write(`Epoch ${epoch + 1}:\n`);
+          samples?.write('Sample (temp=0.3):\n');
+          samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 0.3 })).join(''));
+          samples?.write('\nEnd of sample.\n');
+          samples?.write('Sample (temp=0.7):\n');
+          samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 0.7 })).join(''));
+          samples?.write('\nEnd of sample.\n');
+          samples?.write('Sample (temp=1):\n');
+          samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 1 })).join(''));
+          samples?.write('\nEnd of sample.\n');
+          samples?.write('Sample (temp=1.3):\n');
+          samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 1.3 })).join(''));
+          samples?.write('\nEnd of sample.\n\n');
+          samples?.write('Sample (argmax):\n');
+          samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 'max' })).join(''));
+          samples?.write('\nEnd of sample.\n\n');
+
+          snapshotted = end;
+        }
+      }
+    }
+    logs?.write('Training completed.\n');
+  }
+
+  public async transform(
+    tids: TID[],
+    { temp = 1, logits, llo = false }: {
+      temp?: number | 'max';
+      logits?: Writable;
+      llo?: boolean;
+    } = {},
+  ): Promise<[TID[], TID[]]> {
+    const itids = [
+      ...Array(Math.max(0, this.ws - tids.length)).fill(0),
+      ...tids.slice(-this.ws),
+    ];
+
+    await this.processing.itids.set(itids);
+    await this.processing.ttids.set([...itids.slice(1), 0]);
+    const ilogits = await this.ilogits(this.processing.itids);
+    const e = await this.e(ilogits);
+    const q = this.q(e);
+    const k = this.k(e);
+    const v = this.v(e);
+    const s = this.s(q, k);
+    const wa = await this.wa(s);
+    const a = this.a(wa, v);
+    const o = this.o(a);
+    const ologits = await this.ologits(o, temp === 'max' ? 1 : temp);
+    const out = temp === 'max' ? this.otids(ologits) : this.otids(await this.rnd(), ologits);
+    const otids = await out.get();
+
+    await this.debug(0, this.processing.itids, this.processing.ttids, o, ologits, logits, llo);
+
+    return [itids, otids];
+  }
+
+  public async run(
+    tids: TID[],
+    { temp = 1, samples = this.ws, logits }: {
+      temp?: number | 'max';
+      samples?: number;
+      logits?: Writable;
+    } = {},
+  ): Promise<TID[]> {
+    const result: TID[] = [];
+    while (result.length < samples) {
+      const [_, out] = await this.transform(tids, { temp, logits, llo: true });
+      const next: TID = out[this.ws - 1];
+      result.push(next);
+      tids.shift();
+      tids.push(next);
+    }
+    return result;
+  }
+
+  public async toJSON(): Promise<ModelData<V, W, D>> {
+    return {
+      tokenizer: `${this.tokenizer.tokens.length}${this.tokenizer.tokens[0]}`,
+      vs: this.vs,
+      ws: this.ws,
+      ds: this.ds,
+      stage: this.stage,
+      samples: this.samples,
+      weights: {
+        inp: await this.weights.inp.get(),
+        pos: await this.weights.pos.get(),
+        tqry: await this.weights.tqry.get(),
+        tkey: await this.weights.tkey.get(),
+        tval: await this.weights.tval.get(),
+        out: await this.weights.tout.get(),
+      },
+    };
+  }
+}

src/tokenizer.ts

@@ -0,0 +1,160 @@
+import { GroupQueueBase, QueueItem } from './data/group-queue';
+import { List, ListItem } from './data/list';
+import { Trie } from './data/trie';
+
+interface InternalStringPairQueueItem extends QueueItem<ListItem<string>> {
+  prev?: StringPairQueueItem;
+  next?: StringPairQueueItem;
+}
+
+interface StringPairQueueItem extends QueueItem<ListItem<string>> {
+  readonly prev?: StringPairQueueItem;
+  readonly next?: StringPairQueueItem;
+}
+
+class StringPairQueue extends GroupQueueBase<string, ListItem<string>> {
+  protected items: Map<ListItem<string>, StringPairQueueItem>;
+
+  public constructor(values?: Iterable<ListItem<string>>) {
+    super();
+    this.items = new Map();
+    for (const v of values ?? []) {
+      this.insert(v);
+    }
+  }
+
+  protected override hash(v: ListItem<string>): string {
+    return `${v.value}$$$$${v.next.value}`;
+  }
+
+  override insert(value: ListItem<string>): StringPairQueueItem {
+    const item: StringPairQueueItem = super.insert(value);
+
+    const prev: StringPairQueueItem | undefined = this.items.get(value.prev);
+    const next: StringPairQueueItem | undefined = this.items.get(value.next);
+    (item as InternalStringPairQueueItem).prev = prev;
+    (item as InternalStringPairQueueItem).next = next;
+
+    if (prev) {
+      (prev as InternalStringPairQueueItem).next = item;
+    }
+    if (next) {
+      (next as InternalStringPairQueueItem).prev = item;
+    }
+
+    this.items.set(value, item);
+
+    return item;
+  }
+
+  override remove(item: StringPairQueueItem): boolean {
+    if (!super.remove(item)) {
+      return false;
+    }
+    if (item.prev) {
+      (item.prev as InternalStringPairQueueItem).next = undefined;
+    }
+    if (item.next) {
+      (item.next as InternalStringPairQueueItem).prev = undefined;
+    }
+    this.items.delete(item.value);
+    return true;
+  }
+}
+
+export type TID = number;
+
+export class Tokenizer {
+  protected trie: Trie;
+  protected list: string[];
+
+  public get tokens(): Readonly<string[]> {
+    return this.list;
+  }
+
+  protected train(size: number, text: string): void {
+    if (!text) {
+      return;
+    }
+
+    console.log('Gathering pairs.');
+
+    let log: number = Date.now();
+    const samples: number[] = [];
+
+    const tokens = new List(text);
+    const pairs: StringPairQueue = new StringPairQueue(tokens.items());
+    const vocab = new Set<string>(tokens);
+
+    console.log('Merging pairs.');
+
+    while (vocab.size < size) {
+      const begin: number = Date.now();
+
+      const items: List<StringPairQueueItem> | undefined = pairs.peek();
+      if (!items?.length) {
+        console.log('No merge found.');
+        break;
+      } else {
+        const { value: { value: a, next: { value: b } } } = items.first!;
+        const merged: string = a + b;
+        for (const item of items) {
+          if (!item.orphan) {
+            if (item.prev) {
+              pairs.remove(item.prev);
+            }
+            if (item.next) {
+              pairs.remove(item.next);
+            }
+
+            pairs.remove(item);
+            tokens.set(item.value, merged);
+            tokens.delete(item.value.next);
+
+            pairs.insert(item.value.prev);
+            pairs.insert(item.value);
+          }
+        }
+        vocab.add(merged);
+
+        const end: number = Date.now();
+        samples.push(end - begin);
+        if (end - log > 5000) {
+          console.log(`Merged ${vocab.size} / ${size}, ETA ${Math.round(samples.reduce((a, b) => a + b, 0) / samples.length * (size - vocab.size) / 1000)}s`);
+          log = end;
+        }
+      }
+    }
+    this.trie = new Trie([...vocab].sort((l, r) => r.length - l.length));
+  }
+
+  public constructor(size: number, text: string);
+  public constructor(data: object);
+  public constructor(sizeOrData: number | object, text?: string) {
+    if (text) {
+      this.trie = new Trie();
+      this.train(sizeOrData as number, text!);
+    } else {
+      this.trie = Trie.fromJSON(sizeOrData as object);
+    }
+    this.list = [...this.trie];
+  }
+
+  public tokenize(text: string): TID[] {
+    const tokens: string[] = [];
+    while (text.length) {
+      const match = this.trie.match(text);
+      tokens.push(match ?? text[0]);
+      text = text.slice(match?.length ?? 1);
+    }
+    return tokens.map(t => this.list.indexOf(t));
+  }
+
+  public detokenize(tids: TID[]): string[] {
+    return tids.map(t => this.list[t]);
+  }
+
+  public toJSON(): object {
+    return this.trie.toJSON();
+  }
+}

test/group-queue.spec.ts

@@ -0,0 +1,57 @@
+import { describe, expect, it } from '@jest/globals';
+import { GroupQueue, QueueItem } from '../src/data/group-queue';
+import { List } from '../src/data/list';
+
+describe('GroupQueue', () => {
+  function json<T>(list: List<QueueItem<T>>): T[] {
+    return [...list].map(({ value }) => value);
+  }
+
+  describe('from array', () => {
+    it('[]', () => expect([...new GroupQueue([])].map(json)).toEqual([]));
+    it('[1]', () => expect([...new GroupQueue([1])].map(json)).toEqual([[1]]));
+    it('["b", "c", "b", "c"]', () => expect([...new GroupQueue(['b', 'c', 'b', 'c'])].map(json)).toEqual([['b', 'b'], ['c', 'c']]));
+  });
+
+  describe('length', () => {
+    it('[]', () => expect(new GroupQueue([]).length).toEqual(0));
+    it('[1]', () => expect(new GroupQueue([1]).length).toEqual(1));
+    it('["b", "c", "b", "c"]', () => expect(new GroupQueue(['b', 'c', 'b', 'c']).length).toEqual(2));
+  });
+
+  describe('insert', () => {
+    function insert<T>(queue: GroupQueue<T>, value: T): T[][] {
+      queue.insert(value);
+      return [...queue].map(json);
+    }
+    it('0 -> []', () => expect([...insert(new GroupQueue<number>([]), 0)]).toEqual([[0]]));
+    it('0 -> [1]', () => expect([...insert(new GroupQueue([1]), 0)]).toEqual([[1], [0]]));
+    it('1 -> [1]', () => expect([...insert(new GroupQueue([1]), 1)]).toEqual([[1, 1]]));
+    it('"d" -> ["b", "c", "b", "c"]', () => expect([...insert(new GroupQueue(['b', 'c', 'b', 'c']), 'd')]).toEqual([['b', 'b'], ['c', 'c'], ['d']]));
+    it('"b" -> ["b", "c", "b", "c"]', () => expect([...insert(new GroupQueue(['b', 'c', 'b', 'c']), 'b')]).toEqual([['b', 'b', 'b'], ['c', 'c']]));
+  });
+
+  describe('itemOf / remove', () => {
+    function remove<T>(queue: GroupQueue<T>, value: T): T[][] {
+      const item = queue.itemOf(value);
+      if (item) {
+        queue.remove(item);
+      }
+      return [...queue].map(json);
+    }
+    it('[] - 0', () => expect([...remove(new GroupQueue<number>([]), 0)]).toEqual([]));
+    it('[1] - 0', () => expect([...remove(new GroupQueue([1]), 0)]).toEqual([[1]]));
+    it('[1] - 1', () => expect([...remove(new GroupQueue([1]), 1)]).toEqual([]));
+    it('["b", "c", "b", "c"] - "d"', () => expect([...remove(new GroupQueue(['b', 'c', 'b', 'c']), 'd')]).toEqual([['b', 'b'], ['c', 'c']]));
+    it('["b", "c", "b", "c"] - "b"', () => expect([...remove(new GroupQueue(['b', 'c', 'b', 'c']), 'b')]).toEqual([['c', 'c'], ['b']]));
+  });
+
+  describe('pop', () => {
+    function pop<T>(queue: GroupQueue<T>): [T[], T[][]] {
+      return [[...queue.pop() ?? []].map(({ value }) => value), [...queue].map(json)];
+    }
+    it('[]', () => expect([...pop(new GroupQueue<number>([]))]).toEqual([[], []]));
+    it('[1]', () => expect([...pop(new GroupQueue([1]))]).toEqual([[1], []]));
+    it('["b", "c", "b", "c"]', () => expect([...pop(new GroupQueue(['b', 'c', 'b', 'c']))]).toEqual([['b', 'b'], [['c', 'c']]]));
+  });
+});

test/list.spec.ts

@@ -0,0 +1,88 @@
+import { describe, expect, it } from '@jest/globals';
+import { List, ListItem } from '../src/data/list';
+
+describe('List', () => {
+  describe('from / to array', () => {
+    it('[]', () => expect([...new List([])]).toEqual([]));
+    it('[1]', () => expect([...new List([1])]).toEqual([1]));
+    it('["b", "c", "b", "c"]', () => expect([...new List(['b', 'c', 'b', 'c'])]).toEqual(['b', 'c', 'b', 'c']));
+  });
+
+  describe('length', () => {
+    it('[]', () => expect(new List([]).length).toEqual(0));
+    it('[1]', () => expect(new List([1]).length).toEqual(1));
+    it('["b", "c", "b", "c"]', () => expect(new List(['b', 'c', 'b', 'c']).length).toEqual(4));
+  });
+
+  describe('first', () => {
+    it('[]', () => expect(new List([]).first).toBeUndefined());
+    it('[1]', () => expect(new List([1]).first).toEqual(1));
+    it('["b", "c", "b", "c"]', () => expect(new List(['b', 'c', 'b', 'c']).first).toEqual('b'));
+  });
+
+  describe('last', () => {
+    it('[]', () => expect(new List([]).last).toBeUndefined());
+    it('[1]', () => expect(new List([1]).last).toEqual(1));
+    it('["b", "c", "b", "c"]', () => expect(new List(['b', 'c', 'b', 'c']).last).toEqual('c'));
+  });
+
+  describe('prepend', () => {
+    function prepend<T>(list: List<T>, value: T): List<T> {
+      list.prepend(value);
+      return list;
+    }
+    it('0 -> []', () => expect([...prepend(new List<number>([]), 0)]).toEqual([0]));
+    it('0 -> [1]', () => expect([...prepend(new List([1]), 0)]).toEqual([0, 1]));
+    it('"d" -> ["b", "c", "b", "c"]', () => expect([...prepend(new List(['b', 'c', 'b', 'c']), 'd')]).toEqual(['d', 'b', 'c', 'b', 'c']));
+  });
+
+  describe('append', () => {
+    function append<T>(list: List<T>, value: T): List<T> {
+      list.append(value);
+      return list;
+    }
+    it('[] <- 0', () => expect([...append(new List<number>([]), 0)]).toEqual([0]));
+    it('[1] <- 0', () => expect([...append(new List([1]), 0)]).toEqual([1, 0]));
+    it('["b", "c", "b", "c"] <- "d"', () => expect([...append(new List(['b', 'c', 'b', 'c']), 'd')]).toEqual(['b', 'c', 'b', 'c', 'd']));
+  });
+
+  describe('includes', () => {
+    it('[] not includes 0', () => expect(new List<number>([]).includes(0)).toBe(false));
+    it('[1] not includes 0', () => expect(new List([1]).includes(0)).toBe(false));
+    it('[1] includes 1', () => expect(new List([1]).includes(1)).toBe(true));
+    it('["b", "c", "b", "c"] not includes "d"', () => expect(new List(['b', 'c', 'b', 'c']).includes('d')).toBe(false));
+    it('["b", "c", "b", "c"] includes "b"', () => expect(new List(['b', 'c', 'b', 'c']).includes('b')).toBe(true));
+  });
+
+  describe('itemOf / set', () => {
+    function update<T>(list: List<T>, from: T, to: T): List<T> {
+      const item: ListItem<T> | undefined = list.itemOf(from);
+      if (item) {
+        list.set(item, to);
+      }
+      return list;
+    }
+
+    it('[]: 0 -> 1', () => expect([...update(new List<number>([]), 0, 1)]).toEqual([]));
+    it('[1]: 0 -> 1', () => expect([...update(new List([1]), 0, 1)]).toEqual([1]));
+    it('[1]: 1 -> 0', () => expect([...update(new List([1]), 1, 0)]).toEqual([0]));
+    it('["b", "c", "b", "c"]: "b" -> "d"', () => expect([...update(new List(['b', 'c', 'b', 'c']), 'b', 'd')]).toEqual(['d', 'c', 'b', 'c']));
+    it('["b", "c", "b", "c"]: "d" -> "b"', () => expect([...update(new List(['b', 'c', 'b', 'c']), 'd', 'b')]).toEqual(['b', 'c', 'b', 'c']));
+  });
+
+  describe('itemOf / delete', () => {
+    function remove<T>(list: List<T>, value: T): List<T> {
+      const item: ListItem<T> | undefined = list.itemOf(value);
+      if (item) {
+        list.delete(item);
+      }
+      return list;
+    }
+
+    it('[] - 0', () => expect([...remove(new List<number>([]), 0)]).toEqual([]));
+    it('[1] - 0', () => expect([...remove(new List([1]), 0)]).toEqual([1]));
+    it('[1] - 1', () => expect([...remove(new List([1]), 1)]).toEqual([]));
+    it('["b", "c", "b", "c"] - "b"', () => expect([...remove(new List(['b', 'c', 'b', 'c']), 'b')]).toEqual(['c', 'b', 'c']));
+    it('["b", "c", "b", "c"] - "d"', () => expect([...remove(new List(['b', 'c', 'b', 'c']), 'd')]).toEqual(['b', 'c', 'b', 'c']));
+  });
+});

test/tsconfig.json

@@ -0,0 +1,7 @@
+{
+  "extends": "../tsconfig.json",
+  "include": [
+    "../src/**/*.ts",
+    "**/*.ts",
+  ],
+}

tools/fetch.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+
+model=$1
+count=${2:-9000}
+
+mkdir -p ${model}/downloads
+rm -rf ${model}/downloads/*
+
+BOOK_IDS=(
+  174 1342 11 2701 84 2600 1661 215 1952 98 345 35 1497 2147 4300
+)
+
+echo "Downloading books..."
+
+i=1
+for ID in "${BOOK_IDS[@]}"; do
+  echo "Fetching $ID..."
+  curl -s -L "https://www.gutenberg.org/files/$ID/$ID-0.txt" -o "${model}/downloads/$ID.txt"
+  i=$((i+1))
+  if [ ${i} -gt ${count} ]; then
+    break;
+  fi
+done
+
+echo "Concatenating into corpus.txt..."
+
+# Remove headers/footers and join them
+rm -f ${model}/corpus.txt
+for FILE in ${model}/downloads/*.txt; do
+  awk '/\*\*\* START OF/,/\*\*\* END OF/ { if (!/START|END/) print }' "$FILE" >> ${model}/corpus.txt
+  echo >> ${model}/corpus.txt
+done
+
+echo "All done. Corpus size:"
+du -h ${model}/corpus.txt

tools/stats.py

@@ -0,0 +1,178 @@
+import sys
+import pandas as pd
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+import matplotlib.dates as mdates  # Add this import at the top
+import select
+from datetime import datetime
+import threading
+import queue
+
+fig = plt.figure(figsize=(10, 6))
+
+ax_lr = fig.add_axes([0.05, 0.3, 0.9, 0.65])
+ax_lr.set_facecolor("none")
+ax_lr.set_ylim(bottom=0, top=None)
+ax_lr.yaxis.set_ticks_position('right')
+ax_lr.tick_params(axis='y', labelcolor='tab:brown')
+
+ax_acc = fig.add_axes([0.05, 0.3, 0.9, 0.65], sharex=ax_lr)
+ax_acc.set_facecolor("none")
+ax_acc.set_ylim(0, 100)
+ax_acc.yaxis.set_ticks_position('right')
+ax_acc.tick_params(axis='y', labelcolor='tab:green')
+
+ax_loss = fig.add_axes([0.05, 0.3, 0.9, 0.65], sharex=ax_lr)
+ax_loss.set_facecolor("none")
+ax_loss.set_ylim(bottom=0, top=None)
+ax_loss.tick_params(axis='y', labelcolor='tab:blue')
+
+ax_time = fig.add_axes([0.05, 0.05, 0.9, 0.2])
+ax_time.set_facecolor("none")
+ax_time.set_ylim(bottom=0, top=None)
+ax_time.yaxis.set_ticks_position('right')
+ax_time.tick_params(axis='y', labelcolor='tab:orange')
+
+columns = ["index", "timestamp", "epoch", "batch", "lr", "loss", "accuracy", "grad_inp", "grad_qry", "grad_key", "grad_val", "grad_out", "queue_ms", "compute_ms", "infer_ms"]
+df = pd.DataFrame(columns=columns)
+ema_alpha = 0.1
+window_size = 50
+
+def parse_line(line):
+    try:
+        parts = line.strip().split(",")
+        if parts[0] == "index" or len(parts) < len(columns):
+            return None
+        row = {
+            "index": int(parts[0]),
+            "timestamp": datetime.fromtimestamp(int(parts[1]) / 1000),
+            "epoch": int(parts[2]),
+            "batch": int(parts[3]),
+            "lr": float(parts[4]),
+            "loss": float(parts[5]),
+            "accuracy": float(parts[6]) * 100,
+            "grad_inp": float(parts[7]),
+            "grad_qry": float(parts[8]),
+            "grad_key": float(parts[9]),
+            "grad_val": float(parts[10]),
+            "grad_out": float(parts[11]),
+            "queue_ms": float(parts[12]),
+            "compute_ms": float(parts[13]),
+            "infer_ms": float(parts[14]),
+        }
+        return row
+    except Exception as e:
+        print(f"Bad line: {line.strip()} ({e})", file=sys.stderr)
+        return None
+
+line_queue = queue.Queue()
+
+def stdin_reader():
+    for line in sys.stdin:
+        line_queue.put(line)
+
+threading.Thread(target=stdin_reader, daemon=True).start()
+
+def get_new_lines():
+    lines = []
+    while not line_queue.empty():
+        lines.append(line_queue.get())
+    return lines
+
+last_seen_index = None
+
+def animate(_):
+    global df
+
+    lines = get_new_lines()
+    if not lines:
+        return
+
+    new_rows = [parse_line(line) for line in lines]
+    new_rows = [r for r in new_rows if r is not None]
+
+    if not new_rows:
+      return
+
+    global last_seen_index
+
+    new_index = new_rows[-1]["index"]
+    if last_seen_index and (new_index < last_seen_index):
+        print("Index reset, restarting plot")
+        df = pd.DataFrame(columns=columns)
+
+    last_seen_index = new_index
+
+    df_new = pd.DataFrame(new_rows)
+    df_new["total_ms"] = df_new["queue_ms"] + df_new["compute_ms"] + df_new["infer_ms"]
+    df = pd.concat([df, df_new], ignore_index=True)
+
+    df["ema_loss"] = df["loss"].ewm(alpha=ema_alpha).mean().combine_first(df["loss"])
+    df["avg_loss"] = df["loss"].rolling(window=window_size).mean().combine_first(df["loss"])
+    df["avg_accuracy"] = df["accuracy"].rolling(window=window_size).mean().combine_first(df["accuracy"])
+    df["avg_total_ms"] = df["total_ms"].rolling(window=window_size).mean().combine_first(df["total_ms"])
+
+    ax_lr.clear()
+    ax_lr.set_ylabel("LR")
+    ax_lr.yaxis.set_label_position("right")
+
+    ax_acc.clear()
+    ax_acc.set_ylabel("Accuracy (%)")
+
+    ax_loss.clear()
+    ax_loss.set_title("Live Training Metrics")
+    ax_loss.set_xlabel("Epoch / Batch")
+    ax_loss.set_ylabel("Loss")
+    ax_acc.yaxis.set_label_position("right")
+
+    ax_time.clear()
+    ax_time.set_xlabel("Time")
+    ax_time.set_ylabel("Processing Time (ms)")
+    ax_time.yaxis.set_label_position("right")
+
+    ax_lr.plot(df["index"], df["lr"], label="LR", linestyle='--', alpha=0.3, color="tab:brown")
+
+    ax_acc.plot(df["index"], df["accuracy"], label="Accuracy", linestyle='--', alpha=0.3, color="tab:green")
+    ax_acc.plot(df["index"], df["avg_accuracy"], label="Avg Accuracy",  color="tab:green")
+
+    ax_loss.plot(df["index"], df["loss"], label="Loss", linestyle=':', alpha=0.3, color="tab:blue")
+    ax_loss.plot(df["index"], df["ema_loss"], label="EMA Loss", linestyle='--', alpha=0.7, color="tab:blue")
+    ax_loss.plot(df["index"], df["avg_loss"], label=f"Avg Loss", color="tab:blue")
+
+    ax_time.plot(df["index"], df["queue_ms"], label="Queueing", linestyle=':', alpha=0.3, color="#ffa64d")
+    ax_time.plot(df["index"], df["compute_ms"], label="Computing", linestyle=':', alpha=0.3, color="#cc6600")
+    ax_time.plot(df["index"], df["infer_ms"], label="Inferring", linestyle=':', alpha=0.3, color="#ffb84d")
+    ax_time.plot(df["index"], df["total_ms"], label="Total Time",linestyle='--', alpha=0.7, color="tab:orange")
+    ax_time.plot(df["index"], df["avg_total_ms"], label="Avg Total Time", color="tab:orange")
+
+    # Add epoch/batch label ticks to secondary X axis
+    tick_locs = ax_time.get_xticks()
+    tick_labels = []
+
+    if not df.empty:
+        for loc in tick_locs:
+            idx = (abs(df["index"] - loc)).argmin()
+            e, b = int(df.iloc[idx]["epoch"]), int(df.iloc[idx]["batch"])
+            tick_labels.append(f"{e + 1}/{b + 1}")
+    else:
+        tick_labels = [""] * len(tick_locs)
+
+    ax_loss.set_xlim(ax_time.get_xlim())
+    ax_loss.set_xticks(tick_locs)
+    ax_loss.set_xticklabels(tick_labels)
+
+    # Gather all legend entries
+    lines = []
+    labels = []
+
+    for ax in [ax_lr, ax_acc, ax_loss]:
+        l, lab = ax.get_legend_handles_labels()
+        lines += l
+        labels += lab
+
+    # Add combined legend to ax_loss
+    ax_loss.legend(lines, labels, loc='upper left')  # or 'upper right', etc.
+    ax_time.legend(loc='upper left')
+
+ani = animation.FuncAnimation(fig, animate, interval=1000)
+plt.show()

tsconfig.json

@@ -0,0 +1,24 @@
+{
+  "compilerOptions": {
+    "baseUrl": ".",
+    "outDir": "build",
+    "target": "ES2022",
+    "module": "ES2022",
+    "moduleResolution": "Node",
+    "esModuleInterop": true,
+    "lib": [
+      "DOM", "ES2022"
+    ],
+    "noImplicitAny": true,
+    "noImplicitOverride": true,
+    "noImplicitReturns": true,
+    "noImplicitThis": true,
+    "strictNullChecks": true,
+    "strictPropertyInitialization": true,
+    "declaration": false,
+    "sourceMap": true,
+  },
+  "include": [
+    "src/**/*.ts",
+  ],
+}