# CAX: Cellular Automata Accelerated in JAX

arXiv: https://arxiv.org/abs/2410.02651 · HTML: https://mlsanity.com/paper/2410.02651/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 31 references that are in the map.
- First posted on arXiv 3 Oct 2024.

## Builds on (most cited first, 15 of 15)

- [LLMs and the Abstraction and Reasoning Corpus: Successes, Failures, and the Importance of Object-based Representations](https://mlsanity.com/paper/2305.18354/) (2023)
- [Cellular automata as convolutional neural networks](https://mlsanity.com/paper/1809.02942/) (2018)
- [Lenia - Biology of Artificial Life](https://mlsanity.com/paper/1812.05433/) (2018)
- [evosax: JAX-based Evolution Strategies](https://mlsanity.com/paper/2212.04180/) (2022)
- [Learning Graph Cellular Automata](https://mlsanity.com/paper/2110.14237/) (2021)
- [EvoJAX: Hardware-Accelerated Neuroevolution](https://mlsanity.com/paper/2202.05008/) (2022)
- [Variational Neural Cellular Automata](https://mlsanity.com/paper/2201.12360/) (2022)
- [Attention-based Neural Cellular Automata](https://mlsanity.com/paper/2211.01233/) (2022)
- [Lenia and Expanded Universe](https://mlsanity.com/paper/2005.03742/) (2020)
- [HyperNCA: Growing Developmental Networks with Neural Cellular Automata](https://mlsanity.com/paper/2204.11674/) (2022)
- [DyNCA: Real-time Dynamic Texture Synthesis Using Neural Cellular Automata](https://mlsanity.com/paper/2211.11417/) (2022)
- [QDax: A Library for Quality-Diversity and Population-based Algorithms with Hardware Acceleration](https://mlsanity.com/paper/2308.03665/) (2023)
- [Toward Artificial Open-Ended Evolution within Lenia using Quality-Diversity](https://mlsanity.com/paper/2406.04235/) (2024)
- [Goal-Guided Neural Cellular Automata: Learning to Control Self-Organising Systems](https://mlsanity.com/paper/2205.06806/) (2022)
- [Biomaker CA: a Biome Maker project using Cellular Automata](https://mlsanity.com/paper/2307.09320/) (2023)

## Built on by (most cited first, 8 of 8)

- [ARC-NCA: Towards Developmental Solutions to the Abstraction and Reasoning Corpus](https://mlsanity.com/paper/2505.08778/) (2025)
- [Neural Cellular Automata for ARC-AGI](https://mlsanity.com/paper/2506.15746/) (2025)
- [Neural cellular automata: applications to biology and beyond classical AI](https://mlsanity.com/paper/2509.11131/) (2025)
- [A Path to Universal Neural Cellular Automata](https://mlsanity.com/paper/2505.13058/) (2025)
- [rNCA: Self-Repairing Segmentation Masks](https://mlsanity.com/paper/2512.13397/) (2025)
- [CellARC: Measuring Intelligence with Cellular Automata](https://mlsanity.com/paper/2511.07908/) (2025)
- [Evolving Many Worlds: Towards Open-Ended Discovery in Petri Dish NCA via Population-Based Training](https://mlsanity.com/paper/2604.11248/) (2026)
- [On Growth and Form, and Function: Reusable Regulatory Handles Control Phenotypic Variation](https://mlsanity.com/paper/2609.29755/) (2026)

## Abstract (arXiv)

Cellular automata have become a cornerstone for investigating emergence and self-organization across diverse scientific disciplines. However, the absence of a hardware-accelerated cellular automata library limits the exploration of new research directions, hinders collaboration, and impedes reproducibility. In this work, we introduce CAX (Cellular Automata Accelerated in JAX), a high-performance and flexible open-source library designed to accelerate cellular automata research. CAX delivers cutting-edge performance through hardware acceleration while maintaining flexibility through its modular architecture, intuitive API, and support for both discrete and continuous cellular automata in arbitrary dimensions. We demonstrate CAX's performance and flexibility through a wide range of benchmarks and applications. From classic models like elementary cellular automata and Conway's Game of Life to advanced applications such as growing neural cellular automata and self-classifying MNIST digits, CAX speeds up simulations up to 2,000 times faster. Furthermore, we demonstrate CAX's potential to accelerate research by presenting a collection of three novel cellular automata experiments, each implemented in just a few lines of code thanks to the library's modular architecture. Notably, we show that a simple one-dimensional cellular automaton can outperform GPT-4 on the 1D-ARC challenge.