[Chess] Stockfish-style legal move generation: ~11x faster step, fixes silent move truncation

[gweber]

Summary

This PR replaces the make-move legality filter in chess _legal_action_mask (up to 200x _apply_move + _is_checked per position) with the classic checkers/pins/king-danger scheme used by conventional engines, computed once per position:

• Checker detection around the king (near pieces + sliders with a clear path)
• Check-evasion target mask: non-king moves must capture the single checker or block its line; empty on double check
• Absolute pin rays: the first own piece along each of the 8 king rays, backed by a matching enemy slider, is restricted to moves along that ray (new RAYS / RAY_DIR / IS_DIAG_DIR tables)
• King-danger squares evaluated with the king lifted off the board, so sliders keep attacking through the vacated square

After these masks are built, every pseudo-legal move is decided by table lookups — no per-move make/unmake. En passant (at most two candidates, with its well-known edge cases) keeps the exact make-move test. Castling and underpromotion logic are unchanged. _is_attacked now takes the board array instead of GameState so it can run on the king-removed board.

is_terminal now reuses hash_history[0] (already written by _update_history) instead of recomputing the Zobrist hash.

Fixes a silent correctness bug

The old code compacted candidate moves through jnp.nonzero(..., size=200), silently dropping legal moves in positions with more than 200 pseudo-legal candidates. Example: the known 218-legal-move position

R6R/3Q4/1Q4Q1/4Q3/2Q4Q/Q4Q2/pp1Q4/kBNN1KB1 w - - 0 1

returned exactly 200 legal actions (18 dropped, no error). Multi-queen positions of this kind do occur in self-play RL, where the policy then trains on a wrong action mask. With this PR there is no cap anymore.

Validation

tests/diff_vs_python_chess.py (included, standalone tool — intentionally not test_* so it doesn't extend regular CI): differential test against python-chess as ground truth.

• 17 adversarial FENs (kiwipete, perft positions 3–6, en-passant pins, ep capture of a checking pawn, castling through attacked squares, double check, pinned-knight zugzwang, the 218-move position) with full depth-2 expansion (every child of every legal move compared again)
• 20 random games, every position compared move-by-move

Zero mismatches.

Benchmark

jax.jit(jax.vmap(...)), batch 2048, CUDA (NVIDIA GB10):

	before	after	speedup
_legal_action_mask	7.505 ms	0.535 ms	14.0x
full Game.step	8.303 ms	0.767 ms	10.8x

The new code only uses gather/scatter/where/argmax primitives, so it is also friendly to non-CUDA backends (cf. #1317).

Relation to other work

• Independent of Make Zobrist hashing Apple Metal (jax-metal) compatible #1317 (Zobrist/Metal) — applies directly onto main.
• Touches the same hotspot as the stale memory/speed branches [Chess] reduce memory usage #1258 (chess/bb) and [Chess] Use int8 #1259 (chess/int8); this approach keeps the existing int32 mailbox representation and data layout, so observation/Zobrist/history code is untouched.