Train MNIST with TypeShade compute

This reproducible experiment trains a softmax classifier on handwritten digits. TypeShade executes the model computation; the host loads data and schedules batches.

Execution boundary: The live lab trains a small real-data subset in your browser. The full-dataset results below are from separate reproducible experiments.

Try MNIST in your browser

Train the actual TypeShade compute program, try test digits, then draw your own digit to classify.

TypeShade code running this lab: forward pass

This real compute function transforms 28×28 pixels into ten digit scores. Both training and prediction use the same forward stage.

@compute([64])
export function forward(@builtin("global_invocation_id") gid: vec3u) {
  const row = gid.x;
  if (row >= batch.count) { return; }
  for (let c: u32 = 0; c < 10; c++) {
    let z = bias[c];
    for (let p: u32 = 0; p < 784; p++) {
      z += pixels[(batch.offset + row) * 784 + p] * weights[p * 10 + c];
    }
    logits[row * 10 + c] = z;
  }
}

Explore the seven (7) explained code stages below

View the TypeShade code running this model

This is the exact softmax.shade.ts file compiled into the WebGPU/WebGL2 model above. All six compute entries are shown, including training and prediction.

"use typeshade";

class Batch {
  count: u32;
  offset: u32;
  rate: f32;
}
declare const batch: uniform<Batch>;
declare const pixels: storage<array<f32>>;
declare const labels: storage<array<u32>>;
declare const weights: storage<array<f32>, "read_write">;
declare const bias: storage<array<f32>, "read_write">;
declare const logits: storage<array<f32>, "read_write">;
declare const delta: storage<array<f32>, "read_write">;
declare const losses: storage<array<f32>, "read_write">;
declare const stats: storage<array<f32>, "read_write">;
declare const gradW: storage<array<f32>, "read_write">;
declare const gradB: storage<array<f32>, "read_write">;
declare const probabilities: storage<array<f32>, "read_write">;
declare const predicted: storage<array<u32>, "read_write">;

function stableProbability(z: f32, peak: f32, total: f32): f32 {
  return exp(z - peak) / total;
}

@compute([64])
export function forward(@builtin("global_invocation_id") gid: vec3u) {
  const row = gid.x;
  if (row >= batch.count) { return; }
  for (let c: u32 = 0; c < 10; c++) {
    let z = bias[c];
    for (let p: u32 = 0; p < 784; p++) {
      z += pixels[(batch.offset + row) * 784 + p] * weights[p * 10 + c];
    }
    logits[row * 10 + c] = z;
  }
}

@compute([64])
export function objective(@builtin("global_invocation_id") gid: vec3u) {
  const row = gid.x;
  if (row >= batch.count) { return; }
  let peak = logits[row * 10];
  for (let c: u32 = 1; c < 10; c++) { peak = max(peak, logits[row * 10 + c]); }
  let total: f32 = 0.;
  for (let c: u32 = 0; c < 10; c++) { total += exp(logits[row * 10 + c] - peak); }
  const label = labels[batch.offset + row];
  // Subtract before adding log(total), retaining small losses at large offsets.
  losses[row] = (peak - logits[row * 10 + label]) + log(total);
  for (let c: u32 = 0; c < 10; c++) {
    delta[row * 10 + c] = (stableProbability(logits[row * 10 + c], peak, total) - (c === label ? 1. : 0.)) / f32(batch.count);
  }
}

@compute([1])
export function reduce() {
  let total: f32 = 0.;
  let correct: f32 = 0.;
  for (let row: u32 = 0; row < batch.count; row++) {
    total += losses[row];
    let best: u32 = 0;
    for (let c: u32 = 1; c < 10; c++) {
      if (logits[row * 10 + c] > logits[row * 10 + best]) { best = c; }
    }
    if (best === labels[batch.offset + row]) { correct += 1.; }
  }
  stats[0] = total / f32(batch.count);
  stats[1] = correct;
}

// Explicit backward, not automatic differentiation. Storage-array reverse mode from
// change 0056 is not implemented on the inspected main. One writer per parameter.
@compute([64])
export function backward(@builtin("global_invocation_id") gid: vec3u) {
  const k = gid.x;
  if (k < 7840) {
    const p = k / 10;
    const c = k % 10;
    let sum: f32 = 0.;
    for (let row: u32 = 0; row < batch.count; row++) {
      sum += pixels[(batch.offset + row) * 784 + p] * delta[row * 10 + c];
    }
    gradW[k] = sum;
  }
  if (k < 10) {
    let sum: f32 = 0.;
    for (let row: u32 = 0; row < batch.count; row++) { sum += delta[row * 10 + k]; }
    gradB[k] = sum;
  }
}

@compute([64])
export function update(@builtin("global_invocation_id") gid: vec3u) {
  const k = gid.x;
  if (k < 7840) { weights[k] -= batch.rate * gradW[k]; }
  if (k < 10) { bias[k] -= batch.rate * gradB[k]; }
}

// One label-free inference entry, shared by the browser and the experiment host.
// The first forward row is the digit being classified. All probability math stays
// inside TypeShade, not in a second JavaScript implementation.
@compute([1])
export function predict() {
  let peak = logits[0];
  let best: u32 = 0;
  for (let c: u32 = 1; c < 10; c++) {
    const z = logits[c];
    if (z > peak) { peak = z; best = c; }
  }
  let total: f32 = 0.;
  for (let c: u32 = 0; c < 10; c++) { total += exp(logits[c] - peak); }
  for (let c: u32 = 0; c < 10; c++) {
    probabilities[c] = stableProbability(logits[c], peak, total);
  }
  predicted[0] = best;
}

Open this exact source revision on GitHub

Backend Training images Epochs

Enter a whole number from 1 to 100. One epoch passes through all selected training images.

More epochs repeat GPU work and increase runtime, power use and device heat. Runtime depends on your backend and image count. Higher epochs do not guarantee better test accuracy. Cancel stops after the current batch.

On start, the browser downloads approximately 12 MB of official MNIST gzip files. They are processed locally. No images or drawings are sent to TypeShade servers.

Choose a backend and start training.

The bar counts training images across all epochs. Evaluation runs separately after each epoch.

Watch the model learn

Each training batch predicts digit scores, calculates loss and gradients, then updates weights. After each epoch, evaluation measures the updated model without changing weights.

Choose a backend and start training.

The curves show measured results on the same 64 test images after each epoch. Lower loss means less prediction error; accuracy counts correct classifications. This subset does not represent the full MNIST test set.

Test cross-entropy loss by epoch

Test accuracy by epoch

Draw a digit

Draw a white digit in the square after training. TypeShade forward and predict kernels classify it.

Scores for digits 0 to 9

Training progress

The host loads images and controls the batches. Training and inference logits run through the selected TypeShade WebGPU or WebGL2 backend. WebGL2 uses lowered GLSL ES 3.00 passes.

The training pipeline

The model has 784 input features and 10 output classes. Training and inference use six TypeShade compute entries.

These are excerpts from the exact TypeShade source compiled for the browser lab. This model is softmax regression: it receives 28×28 pixels (784 inputs) and produces 10 scores for digits 0 through 9. It is neither a CNN nor a pretrained model.

Training: forward → objective → backward → update

Evaluation: forward → objective → reduce

Handwritten inference: forward → predict

Read what each stage computes and why, then inspect the real source directly underneath. Every code sample is extracted at build time from the compiled softmax.shade.ts file.

1. Declare GPU resources in TypeScript

A "use typeshade" file declares pixels, labels, weights, biases, gradients, and output buffers. The batch uniform carries the current mini-batch size and learning rate.

"use typeshade";

class Batch {
  count: u32;
  offset: u32;
  rate: f32;
}
declare const batch: uniform<Batch>;
declare const pixels: storage<array<f32>>;
declare const labels: storage<array<u32>>;
declare const weights: storage<array<f32>, "read_write">;
declare const bias: storage<array<f32>, "read_write">;
declare const logits: storage<array<f32>, "read_write">;
declare const delta: storage<array<f32>, "read_write">;
declare const losses: storage<array<f32>, "read_write">;
declare const stats: storage<array<f32>, "read_write">;
declare const gradW: storage<array<f32>, "read_write">;
declare const gradB: storage<array<f32>, "read_write">;
declare const probabilities: storage<array<f32>, "read_write">;
declare const predicted: storage<array<u32>, "read_write">;

function stableProbability(z: f32, peak: f32, total: f32): f32 {
  return exp(z - peak) / total;
}

An image contains 784 f32 pixels; 10 classes require 784×10 = 7,840 weights. Weights and biases remain GPU-resident and are reused by the next batch. stableProbability is a shared helper for numerically stable softmax.

2. Forward pass: turn pixels into class scores

forward multiplies 784 pixel values by class-specific weights, sums them and adds biases, producing 10 logits per image.

@compute([64])
export function forward(@builtin("global_invocation_id") gid: vec3u) {
  const row = gid.x;
  if (row >= batch.count) { return; }
  for (let c: u32 = 0; c < 10; c++) {
    let z = bias[c];
    for (let p: u32 = 0; p < 784; p++) {
      z += pixels[(batch.offset + row) * 784 + p] * weights[p * 10 + c];
    }
    logits[row * 10 + c] = z;
  }
}

gid.x identifies the image row. For class c, it computes z = bias[c] + sum pixels[p]×weights[p×10+c]. @compute([64]) specifies the workgroup size, with an explicit guard for short batches.

3. Loss and error signal: measure the prediction

objective uses stable softmax and cross-entropy, subtracting the maximum logit before exponentiation.

@compute([64])
export function objective(@builtin("global_invocation_id") gid: vec3u) {
  const row = gid.x;
  if (row >= batch.count) { return; }
  let peak = logits[row * 10];
  for (let c: u32 = 1; c < 10; c++) { peak = max(peak, logits[row * 10 + c]); }
  let total: f32 = 0.;
  for (let c: u32 = 0; c < 10; c++) { total += exp(logits[row * 10 + c] - peak); }
  const label = labels[batch.offset + row];
  // Subtract before adding log(total), retaining small losses at large offsets.
  losses[row] = (peak - logits[row * 10 + label]) + log(total);
  for (let c: u32 = 0; c < 10; c++) {
    delta[row * 10 + c] = (stableProbability(logits[row * 10 + c], peak, total) - (c === label ? 1. : 0.)) / f32(batch.count);
  }
}

The delta value is (predicted probability - one-hot label) / batch size. backward uses this error signal to determine each weight gradient. Both exp and log execute in the TypeShade compute program.

4. Backward pass: accumulate weight gradients

backward computes gradients for 7,840 weights and 10 biases using the batch pixels and delta values.

@compute([64])
export function backward(@builtin("global_invocation_id") gid: vec3u) {
  const k = gid.x;
  if (k < 7840) {
    const p = k / 10;
    const c = k % 10;
    let sum: f32 = 0.;
    for (let row: u32 = 0; row < batch.count; row++) {
      sum += pixels[(batch.offset + row) * 784 + p] * delta[row * 10 + c];
    }
    gradW[k] = sum;
  }
  if (k < 10) {
    let sum: f32 = 0.;
    for (let row: u32 = 0; row < batch.count; row++) { sum += delta[row * 10 + k]; }
    gradB[k] = sum;
  }
}

Each gradW[p×10+c] is sum pixels[row,p]×delta[row,c], while each gradB is a class-wise delta sum. This implementation uses explicitly authored gradients, not automatic differentiation.

5. SGD: update the GPU-resident model

update multiplies each gradient by the learning rate and subtracts it from the corresponding weight or bias.

@compute([64])
export function update(@builtin("global_invocation_id") gid: vec3u) {
  const k = gid.x;
  if (k < 7840) { weights[k] -= batch.rate * gradW[k]; }
  if (k < 10) { bias[k] -= batch.rate * gradB[k]; }
}

// One label-free inference entry, shared by the browser and the experiment host.
// The first forward row is the digit being classified. All probability math stays
// inside TypeShade, not in a second JavaScript implementation.

weights[k] -= batch.rate × gradW[k] is the key step. JavaScript does not read the model back and update it for every batch. The next forward pass observes the GPU-resident updated values.

6. Evaluation: aggregate loss and correct answers

reduce calculates mean cross-entropy and the number of correctly classified examples in a batch.

@compute([1])
export function reduce() {
  let total: f32 = 0.;
  let correct: f32 = 0.;
  for (let row: u32 = 0; row < batch.count; row++) {
    total += losses[row];
    let best: u32 = 0;
    for (let c: u32 = 1; c < 10; c++) {
      if (logits[row * 10 + c] > logits[row * 10 + best]) { best = c; }
    }
    if (best === labels[batch.offset + row]) { correct += 1.; }
  }
  stats[0] = total / f32(batch.count);
  stats[1] = correct;
}

// Explicit backward, not automatic differentiation. Storage-array reverse mode from
// change 0056 is not implemented on the inspected main. One writer per parameter.

stats[0] is mean loss and stats[1] is the correct count. This is a metric reduction, not weight-gradient reduction. The host reads this small output to compute evaluation totals.

7. Prediction: calculate digit 0–9 probabilities on GPU

predict normalizes 10 forward logits and picks the highest-scoring digit without needing labels.

@compute([1])
export function predict() {
  let peak = logits[0];
  let best: u32 = 0;
  for (let c: u32 = 1; c < 10; c++) {
    const z = logits[c];
    if (z > peak) { peak = z; best = c; }
  }
  let total: f32 = 0.;
  for (let c: u32 = 0; c < 10; c++) { total += exp(logits[c] - peak); }
  for (let c: u32 = 0; c < 10; c++) {
    probabilities[c] = stableProbability(logits[c], peak, total);
  }
  predicted[0] = best;
}

The compute entry writes probabilities[c] and predicted[0]. Browser JavaScript only renders those GPU results; it does not run a second Math.exp softmax implementation.

How does the browser connect these TypeShade kernels?

The JavaScript host downloads official IDX files, normalizes pixels to 0–1, and dispatches the TypeShade compute stages. The methods below come directly from the shared browser session, not from illustrative pseudocode.

Actual code for one training batch

trainBatch(pixels, labels, rate) {
      return queue(async () => {
        if (!Number.isFinite(rate) || rate <= 0) throw new RangeError('Invalid learning rate');
        setBatch(pixels, labels, labels.length);
        for (const entry of ['forward', 'objective', 'backward', 'update'] as const)
          await dispatch(entry, labels.length, rate);
      });
    },

setBatch writes pixels and labels to GPU input buffers. The session dispatches forward → objective → backward → update in order, retaining the same GPU-resident weights for subsequent batches.

Actual code for classifying a drawn digit

predict(pixels) {
      return queue(async () => {
        if (pixels.length !== INPUTS) throw new RangeError('MNIST inference requires 784 pixels');
        setBatch(pixels, new Uint32Array(1), 1);
        await dispatch('forward', 1, 0);
        await dispatch('predict', 1, 0);
        const probabilities = (await data.probabilities.read()) as Float32Array;
        const prediction = (await data.predicted.read()) as Uint32Array;
        const sum = probabilities.reduce((a, b) => a + b, 0);
        if (
          !probabilities.every((value) => Number.isFinite(value) && value >= 0) ||
          Math.abs(sum - 1) > 0.0002 ||
          prediction[0] >= CLASSES
        )
          throw new Error('Invalid TypeShade inference probability output');
        return { probabilities, predicted: prediction[0] };
      });
    },

forward → predict computes probabilities and the top class on the GPU. Only 10 floats and one class index are read back. Training and prediction share a serialized queue to avoid races over mutable buffers.

Execution boundary: TypeShade owns the model math (forward, softmax, loss, gradients, SGD and inference). JavaScript handles downloads, batch scheduling, canvas input, UI updates and aggregating metrics. WebGL2 uses TypeShade-generated GLSL ES 3.00 passes, not native compute.

Validated results

The full-data rows come from the committed result record. The WebGL2 validation is a separate, smaller execution gate.

Execution backend	Train / test, epochs	Test accuracy
Generated CPU	60,000 / 10,000, 5 epochs	91.67%
WebGPU (SwiftShader)	60,000 / 10,000, 5 epochs	91.67%
WebGL2 passes	1,024 train / 1,000 test subset	Full-data result not established

The measured WebGPU adapter was Google SwiftShader, a software renderer; these results do not demonstrate physical GPU acceleration.

The real WebGL2 test executed every pass, checked against an independent numeric reference and completed subset training. Evidence: MNIST WebGL2 validation workflow.

Reproduce the experiment

From the compiler repository root with Node 24 and Bun, install dependencies and run the strict WebGL2 tests:

bun install
bun run build
./node_modules/.bin/playwright install --only-shell chromium
node --experimental-strip-types journeys/mnist/run.ts --download --train 1 --test 1 --epochs 1
node --experimental-strip-types journeys/mnist/test-webgl2.mjs --mnist --count 1024 --test 1000 --epochs 5

For the full dataset using the WebGPU backend after downloading the dataset:

node --experimental-strip-types journeys/mnist/run.ts --tier webgpu

The commands use the local Chromium runner and download MNIST data on demand. Performance depends on the actual adapter and host environment.

What this does and does not verify

WebGL2 uses generated GLSL ES 3.00 passes, not native compute shaders.
The WebGL2 CI validates numeric stages and a small real-data training subset; full-dataset WebGL2 accuracy and speed remain unverified.
The published full-dataset WebGPU result used software rendering. No physical GPU speedup is claimed.
The live lab executes TypeShade kernels in the visitor’s browser on a selected subset; full-dataset results were recorded separately.

Source and verification

The source and results are bound to the compiler commit used to build this site.

MNIST in TypeNN

This guide trains a softmax regression with kernels written by hand. TypeNN is a neural-network library built on TypeShade. A model is a stack of layers, and each layer runs as TypeShade kernel functions. The code below defines a model with one hidden layer and trains it on batches.

import * as nn from "typenn";

const device = await nn.device();
const model = new nn.Sequential(
  new nn.Flatten(),
  new nn.Linear(784, 128),
  new nn.ReLU(),
  new nn.Dropout(0.2),
  new nn.Linear(128, 10),
).to(device);
const optimizer = nn.optim.Adam(model.parameters(), { lr: 0.001 });
const criterion = new nn.CrossEntropyLoss();

for await (const batch of batches) {
  await optimizer.zeroGrad();
  const images = device.tensor(batch.images, {
    shape: [batch.size, 1, 28, 28],
  });
  const loss = criterion.forward(model.forward(images), batch.labels);
  await loss.backward();
  await optimizer.step();
  await loss.disposeGraph(model.parameters());
}

Edit this page Report a problem