cemoss17/maia3-onnx

maia3-js

JavaScript inference for Maia3 (Chessformer) — predicts human-like chess moves and per-move WDL evaluations.

ESM-only. Node + browser. ONNX Runtime under the hood.

Install

# Node
npm install maia3-js onnxruntime-node

# Browser
npm install maia3-js onnxruntime-web

onnxruntime-node and onnxruntime-web are optional peer dependencies — install whichever runtime you need.

Quickstart (Node)

import { Maia3, STARTPOS_FEN } from "maia3-js";

const maia = new Maia3({ variant: "5m" });
await maia.load();

const result = await maia.predict({
  fen: STARTPOS_FEN,
  selfElo: 1500,
});

console.log(result.bestMove);          // e.g. "e2e4"
console.log(result.winProbability);    // e.g. 0.52
for (const c of result.candidates) {
  console.log(c.uci, c.probability.toFixed(3), c.wdl);
}

Quickstart (Browser)

import { Maia3, STARTPOS_FEN } from "maia3-js/web";

const maia = new Maia3({
  variant: "5m",
  onProgress: (loaded, total) => console.log(loaded / total),
});
await maia.load();

const r = await maia.predict({ fen: STARTPOS_FEN, selfElo: 1500 });

5m is bundled with the npm package (~21 MB). 23m and 79m are downloaded on first use from Hugging Face and cached locally (Node) or in Cache Storage (browser).

History

Maia3 conditions on up to 8 prior board positions. Supply this via one of two equivalent inputs:

// Option A: explicit prior FENs (oldest first, exclude the current position)
await maia.predict({
  fen: currentFen,
  priorFens: [prevFen, prevPrevFen, ...],
  selfElo: 1500,
});

// Option B: UCI moves applied to a base FEN (defaults to startpos)
await maia.predict({
  fen: currentFen,
  priorMoves: ["e2e4", "e7e5", "g1f3", "b8c6"],
  selfElo: 1500,
});

When no history is supplied the current position is used and silently padded to 8 — matches the default UCI engine behavior in the Python reference (use_uci_history=False).

Passing both priorFens and priorMoves throws.

API

class Maia3 {
  constructor(options?: {
    variant?: "5m" | "23m" | "79m";   // default "5m"
    modelPath?: string;                // node: explicit path
    modelDir?: string;                 // node: cache dir
    url?: string;                      // web: override URL
    modelBytes?: ArrayBuffer;          // web: pre-fetched bytes
    onProgress?: (loaded: number, total: number) => void; // web
    executionProviders?: string[];     // pass-through to ORT
    temperature?: number;              // default 0 (argmax bestMove)
    topP?: number;                     // default 1.0
    topK?: number;                     // default 5
  });

  load(): Promise<void>;
  isLoaded(): boolean;
  predict(input: PredictInput): Promise<PredictionResult>;
  close(): Promise<void>;
}

interface PredictInput {
  fen: string;
  selfElo: number;
  oppoElo?: number;        // defaults to selfElo
  priorFens?: string[];    // exclusive with priorMoves
  priorMoves?: string[];
  baseFen?: string;        // for priorMoves; defaults to startpos
  temperature?: number;
  topP?: number;
  topK?: number;
}

interface PredictionResult {
  bestMove: string;                   // UCI of selected move
  moves: Map<string, number>;         // legal UCI → policy probability
  candidates: MoveCandidate[];        // top-K with per-move WDL
  wdl: { win: number; draw: number; loss: number }; // side-to-move POV
  winProbability: number;             // wdl.win + 0.5 * wdl.draw
}

interface MoveCandidate {
  uci: string;
  probability: number;
  wdl: { win: number; draw: number; loss: number };
  winProbability: number;
}

Sampling

• temperature === 0 (default): bestMove is argmax of the masked policy.
• temperature > 0: bestMove is sampled from the policy with optional nucleus filter (topP < 1). Matches the Python sample_from_logits routine.

Per-candidate WDL

For each of the top-topK candidates, maia3-js runs an additional forward pass with swapped Elos to score the position after the candidate move. WDL is reported in the original player's frame (inverted from the post-move side-to-move).

Variants

Alias	Repo	Parameters	Bundled
5m	UofTCSSLab/Maia3-5M	~5 M	yes
23m	UofTCSSLab/Maia3-23M	~23 M	no
79m	UofTCSSLab/Maia3-79M	~79 M	no

ONNX checkpoints are hosted at huggingface.co/cemoss17/maia3-onnx.

Parity with Python

Test fixtures are generated from the official Maia3 PyTorch implementation. The JS port matches it within:

• Encoding: bit-exact (1e-6) for the (64, 96) token tensor.
• Raw logits: 1e-3 absolute on Top-20 logits.
• Policy probabilities: 5% absolute.
• WDL: 5% absolute (current position and per-candidate).

npm run validate exercises this against 13 curated positions and 200 random parity positions in one go.

Re-exporting weights

If you need to export a custom checkpoint or the 23m / 79m variants yourself, install the maia3 Python package and run:

.venv/bin/python scripts/export-onnx.py \
  --variant 5m \
  --output models/maia3_5m.onnx \
  --validate

The --validate flag runs a 32-position parity check between the PyTorch model and the exported ONNX (atol=1e-4).

Limitations

• inferBatch currently runs sequentially under the hood — the dynamo exporter baked a single-batch reshape into the graph. K=5 candidate scoring is still fast (~5×10ms on CPU), but true batched inference will land in a future export.
• Sampling uses Math.random() — no seed parameter in v0.1.
• Only the policy / value heads are exported. The Maia3 ponder head (predicted clock time) is dropped.

License

MIT.