cstr/chessmamba-onnx

ChessMamba, ONNX (single recurrent step)

An ONNX export of TobiasLogic/chessmamba (MIT), a 16.8M-parameter selective state-space model (Mamba/S6) that reads a chess game as a sequence of moves and predicts the next move, promotion and a value. All credit for the model and its training goes to its author; this repository only converts it.

Two files, same weights:

• chessmamba_step_batch.onnx (use this): batch dimension b on every input and output (state is [10, b, 768, 16]), so a search can expand all of a node's candidate moves in one call. With onnxruntime a batch of 10 costs about 1.4 single steps.
• chessmamba_step.onnx: the same graph fixed at batch 1.

Each graph is one incremental step: feed one move and the recurrent state, get the policy for the next move and the updated state. Playing a game costs one step per move regardless of its length. Used by CrispChess.

Inputs

name	type	shape	meaning
from_sq	int64	[b]	from square, a1 = 0 .. h8 = 63 (python-chess numbering)
to_sq	int64	[b]	to square
promo	int64	[b]	0 none, 1 queen, 2 rook, 3 bishop, 4 knight
ply	int64	[b]	1-based index of this move; clamped to 96 inside
is_start	float32	[b, 1]	1 for the very first step (no move, uses the start token), else 0
state	float32	[10, b, 768, 16]	recurrent state; zeros for the first step

Outputs

name	shape	meaning
policy	[b, 4096]	next-move logits, index from * 64 + to
promo_logits	[b, 5]	promotion logits
value	[b, 1]	expected result for the side to move, -1..1 (tanh)
new_state	[10, b, 768, 16]	state to pass to the next step

Start with is_start = 1 and a zero state, then one step per move with is_start = 0.

Verification

Opset 17, IR 8, 23 standard ops. Compared with the PyTorch model's own step_move over a 30-ply game: max absolute policy-logit difference 1.05e-5 (onnxruntime). The batched graph matches the single-step one to 7e-4 over ten children of the same position. About 15-18 ms per step on one CPU thread with onnxruntime. Exported and checked by export_chess_lms.py.

sha256 of chessmamba_step_batch.onnx: b956c1580c719c41bf530c6bb25d766eae7db7f7fef1ad8b033890dd9317bedd; of chessmamba_step.onnx: 03e79ddea3a67d77a4c5b1deb37b5952107baf48d6e5b2ef12693b75339b06f8