Hello HN,I built a complete, game-agnostic implementation of AlphaZero in JAX.- repo: https://github.com/wtedw/nanoAlphaZero - demo (NN + MCTS run locally in your browser): https://nanoalphazero.wtedw.comIt uses no human data, can train grandmaster-level chess models, and supports a variety of games:- Chess - Go 3x3 - 9x9 - Hex 4x4 - 9x9 - Connect Four - Tic Tac ToeYou can also use this repo to train AlphaZero on any custom 2-player, perfect-information game (this needs more documentation).How does it work?At a high level, the entire AlphaZero algorithm gets compiled into a single jitted run_fn that repeatedly performs self-play and model updates: state = make_alphazero() def run_fn(state): games = selfplay(state) # using Gumbel MuZero # Move active games into the self-play buffer # Move completed games into the replay buffer state = train(state, replay_buffer.sample()) return state while True: state = run_fn(state) There are no threads, queues, or distributed workers to manage. It is just one large JAX function.The repo is primarily focused on making large-scale AlphaZero experimentation fast and easy to run. It s optimized for speed and memory efficiency while remaining compact and hackable. Training strong models is secondary and mostly serves as a sanity check that the underlying logic is sound.Documentation on training custom / complex games is sparse, so if you have any questions feel free to message me.