CAX: Cellular Automata Accelerated in JAX - Faldor, Maxence; Cully, Antoine (International Conference on Learning Representations, 2025)