script-doctor. Code for PuzzleJAX, a benchmark for reasoning and learning, that reimplements PuzzleScript, a concise and expressive DSL and game engine for grid-based puzzle games, in JAX.