Accidentally quadratic: buffer copies made MCTS in DeepMind's mctx 3× slower

Article automatically generated from technical news.

I'm training an AlphaZero-style agent (Gumbel MuZero via DeepMind's mctx ) to lay out working factory modules for Factorio: the network places machines, belts and inserters on a grid, and the reward comes from an exact throughput verifier. Everything runs in JAX on a single RTX 5070: 128 environments in one batch, an action space of A = 1729 (3 entity types × 144 cells × 4 rotations + "done"), and a small 474k-parameter conv net in bf16. While benchmarking training configurations