CAX: Cellular Automata Accelerated in JAX
Authors: Maxence Faldor
Antoine Cully
Venue: ICLR 2025 Oral (2025)
Keywords: cellular automata
emergence
self-organization
neural cellular automata
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: CAX is a high-performance and flexible open-source library designed to accelerate cellular automata research.

Abstract: 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.

Published as a conference paper at ICLR 2025
CAX: CELLULAR AUTOMATA ACCELERATED IN JAX
Maxence Faldor
Antoine Cully
Department of Computing
Department of Computing
Imperial College London
Imperial College London London, United Kingdom
London, United Kingdom
m. faldor222 imperialll... carrriialaaao ac. uk
a culll 1 imperiiall..... ac uk
ABSTRACT
Cellular automata have become a cornersione for investigating emergence and self seff-organzzttion across diverse scientific disciplines. However, the absence of a hardware-acceleree cellular automata library limits the exploration of new re- search 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 bench- marks and applications. From classic models like elementary cellular automata and Conway's ' Game of Life to advanced applications such as growing neural cel- lular 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 few lines of code thanks to the library's 's modular archi- tecture. Notably, we show that a simple one-dimensional cellular automaton can outperform GPT-4 on the ID-ARC challenge.
INTRODUCTION
Emergence is a fundamental concept that has captivated thinkers across various fields of human inquiry, including philosophy, science and art (Holland) 2000). This fascinating phenomenon oc- curs when a complex entity exhibits properties that its constituent parts do not possess individually. From the collective intelligence of ant colonies to the formation of snowflakes, self-organization and emergence manifest in myriad ways. The study of seff-organization and emergence holds the promise to unravel deep mysteries, from the origin of life to the development of conciousness.
Cellular Automata (CA) are models of computation that exemplify how complex patterns and so- phisticated behaviors arise from simple components interacting through basic rules. Originating from the work Ulam and von Neumann in the 1940s (Neumann & Burks 1966). these systems gained prominence with Conway's Game of Life in the 1970s (Gardner 1970) and Wolfram's sys- tematic studies in the 1980s (Wolfram) 2002). The discovery that even elementary cellular automata can be Turing-complete underscores their expressiveness (Cook) 2004). CAs serve as a powerful abstraction for investigating self oeff-oraanzzation and emergence, offering insights into complex phe- nomena across scientific domains, from physies and biology to computer science and artificial life.
In recent years, the integration of machine learning techniques with cellular automata has opened new avenues for research morphogenesis (Mordvintsev al., (2220) and self-organization (Ran- dazzo et al. 2020 2021 Najarro et al., (022). Neural Cellular Automata (NCA) represent key ad- vancement in this field, using differentiable neural networks to learn update rules through gradient- based optimization rather than relying manual design. Recent works have combined NCA with advanced architectures like Convolutional Neural Networks (Pajouheshgar et al. 2023). Graph Neu- ral Networks (Grattarola et al. 2021) and Vision Transformers (Tesfaldet et al. 2022 The conver- gence of machine learning and cellular automata research not only deepened our understanding of complex systems but also underscored the growing computational demands of CA experiments.

---

Published as a conference paper at ICLR 2025
Discrete Cellular Automata
Continuous Cellular Automata
Neural Cellular Automata
Figure 1: Cellular Automata types supported in CAX.
Despite their conceptual simplicity, cellular automata simulations can be computationally intensive, especially when scaling to higher dimensions with large numbers of cells or implementing backprop- agation through time for NCAs. Moreover, the implementation of CA research settings has often been fragmented, with individual researchers frequently reimplementing basic functionalities, cre- ating custom implementations across various deep learning frameworks such as TensorFlow, JAX, and PyTorch. As the field continues to grow and attract increasing interest, there is pressing need for a unified, robust library that facilitates collaboration, reproducibility, fast experimentation and exploration of new research directions.
In response to these challenges and opportunities, we present CAX: Cellular Automata Accelerated in JAX, an open-source library with cutting-edge performance. designed to provide flexible and efficient framework for cellular automata research. CAX built on JAX (Bradbury et al.) 2018), regh-performance numerical computing library, enabling to speed up cellular automata simula- tions through massive parallelization across various hardware accelerators such as CPUs, GPUs, and TPUs. CAX flexible and supports both discrete and continuous cellular automata with any num- ber of dimensions, accommodating classic models like elementary cellular automata and Conway's s Game of Life, as well modern variants such as Lenia and Neural Cellular Automata (Figure !
JAX offers efficient vectorization of CA rules, enabling millions of cell updates to be processed simultaneously. It also provides automatic differentiation capabilities to backpropagate through time efficiently, facilitating the training of Neural Cellular Automata. CAX can run experiments with millions cell updates in minutes, reducing computation times by up to 2,000 times compared to traditional implementations in our benchmark. This performance boost opens up new possibilities for large-scale CA experiments that were previously computationally prohibitive.
CAX's flexibility and potential to accelerate research is showcased through three novel cellular automata experiments. Thanks to CAX's modular architecture, each of these experiments is imple- mented in just few lines of code, significantly reducing the barrier to entry for cellular automata research. Notably, we show that a simple one-dimensional cellular automaton implemented with CAX outperforms GPT-4 on the 1D-ARC challenge (Xu al., (202),, see Section [5.]] Finally, to support users and facilitate adoption, CAX comes with high-quality, diverse examples and compre- hensive documentation. The list of implemented CAs is detailed in Table [
BACKGROUND
2.1 CELLULAR AUTOMATA
A cellular automaton is simple model of computation consisting of a regular grid of cells, each in a particular state. The grid can be in any finite number of dimensions. For each cell, set of cells called its neighborhood is defined relative to the specified cell. The grid is updated at discrete time steps according to fixed rule that determines the new state of each cell based on its current state and the states of the cells in its neighborhood.
CA defined by a tuple (C, S,N, Φ), where L is the d-dimensional lattice or grid with channels, S is the cell state set, N C C is the neighborhood of the origin, and SN - S is the local rule. mapping from the grid to the cell state set L + S called a configuration or pattern. In this work, we will simply refer to it by the state of the CA. S(x) represents the state of cell x E L. Additionally, we denote the neighborhood of a cell x E L by N₂ = {x n, n. E N}. and S(N) (S(n), E Nx).
2

---

Published as a conference paper at ICLR 2025
The global rule 4 S S applies the local rule uniformly to all cells in the lattice and is defined such that, for all x L, (x) = ((S(N))) A cellular automaton is initialized with a state So- Then, the state is updated according to the global rule @ at each discrete time step t N. to give,
S (Sa)), S2 = (SS))....
The close connection between CA and recurrent convolutional neural networks has been observed by numerous researchers (Gilpin 2019 Wulff & Hertz 1992 Mordvintsev al. [020 Chan 2020). For example, the general NCA architecture introduced by Mordvintsev et al. (2020) can be concep- tualized as a "recurrent residual convolutional neural network with per-cell dropout".
2.2 CONTROLLABLE CELLULAR AUTOMATA
A controllable cellular automaton (CCA) is generalization of CA that incorporates the ability to accept external inputs at each time step. CCAs formalize the concept of Goal-Guided NCA that has been introduced in the literature by Sudhakaran et al. (2022). The external inputs can modify the behavior of CCAs, offering the possibility to respond dynamically to changing conditions or control signals while maintaining the fundamental principles of cellular automata.
A CCA is defined by a tuple (, , 6), where I is the input set and é SN x INN - S is the controllable local rule. A mapping from the grid to the input set I C + I is called the input. I(x) represents the input of cell x E L. Similarly to the state, we denote I(Nx)
(I(n), E Nx).
The controllable global rule + S I - S defined such that, for all x in L, I), ((NN),, I(Nx)). A controllable cellular automaton is initialized with an initial state So. Then, the state updated according to the controllable global rule and sequence of input (I,)2o each discrete time step E N, to give,
S = (Sp, Io), S2 = ())
As discussed in Section 2.1 CAs can be conceptualized as recurrent convolutional neural networks. However, traditional CAs lack the ability to take external inputs at each time step. CCAs extend the capabilities of traditional CAs by making them responsive to external inputs, akin to recurrent neural networks processing sequential data. CCAs bridge the gap between recurrent convolutional neural networks and cellular automata, opening up new possibilities for modeling complex systems that exhibit both autonomous emergent behavior and responsiveness external control.
2.3 RELATED WORK
The field of CA has spawned numerous tools and libraries to support research and experimentation, with CellPyLib (Antunes 202T) emerging as one of the most popular and versatile options. This Python library offers simple yet powerful interface for working with 1- and 2-dimensional CA, supporting both discrete and continuous states, making it an ideal baseline for comparative studies and further development. While it provides implementations of classic CA models like Conway's and does not support the train-
Game of Life and Wireworld, CellPyLib is not ing of neural cellular automata. Golly is cross-platform application for exploring Conway's Game of Life and many other types of cellular automata. Golly's ' features include 3D CA rules, custom rule loading, and scripting via Lua or Python. While powerful and versatile for traditional CA, Golly is not designed for hardware acceleration or integration with modern machine learning frameworks.
The recent surge in artificial intelligence has increased the availability of computational resources, and encouraged the development of sophisticated tools such as JAX (Bradbury et al., 2018). a high- performance numerical computing library with automatic differentiation and JIT compilation. A rich ecosystem of specialized libraries has emerged around JAX, such Flax (Heek et al., (024) for neural networks, RLax (DecpMind et al. 2020) for reinforcement learning, and EvoSax (Lange, 2022). EvoJax (Tang et al. 2022) and QDax (Chalumeau et al. 2023) for evolutionary algorithms.
In the realm of cellular automata, there have been efforts to implement specific CA models using JAX. For instance. EvoJax (Tang et al.) 2022) and Leniax (Giraud 2022) both provide hardware- accelerated Lenia implementation. Biomaker CA (Randazzo & Mordvintsev (2023), specific CA model focusing biological pattern formation, further demonstrates the potential of JAX in CA
3

---

Published as a conference paper at ICLR 2025
Parcelve
Update
Neighborhoed
Percaption
Vecbor
Isput
Figure 2: High-level architecture of CAX, illustrating the modular design with perceive and update modules. This flexible structure supports various CA types across multiple dimensions. (Adapted from Mordvintsev et al. (2020) under CC-BY 4.0 license.)
research. Finally, various GitHub repositories replicate results from neural cellular automata pa- pers, but these implementations are typically narrow focus. Recent advancements in continuous cellular automata research have also benefited from JAX-based implementations. These include Lenia (Chan 2020) and Leniabreeder (Faldor Cully 2024). which have enabled large-scale sim- ulations of open-ended evolution in continuous cellular automata (Chan) 2023).
While existing implementations demonstrate JAX's potential in CA research, they also reveal sig- nificant gaps in the field. Current tools are often specialized for specific CA types (e.g., discrete, 1- and 2-dimensional), narrow focus (e.g., replicating specifie neural CA papers), or lack hardware acceleration. This limitation underscores the need for comprchensive, flexible, and efficient library that can handle a broad spectrum of CA types while leveraging hardware acceleration. CAX aims to address this gap by providing versatile, JAX-based tool to accelerate progress across the entire landscape of cellular automata research.
3 CAX: CELLULAR AUTOMATA ACCELERATED IN JAX
CAX gegh-performance and flexible open-source library designed to accelerate cellular automata research. In this section, detail CAX's architecture, design and key features. At its core, CAX leverages JAX and Flax (Heek et al. 2024), capitalizing on the well-established connection between CA and recurrent convolutional neural networks. This synergy, discussed in Section 2 allows CAX to harness adancementts in machine learning to accelerate CA research. CAX offers modular and intuitive design through a user- user-friendly interface, supporting both discrete and continuous cellular automata across any number of dimensions. This flexibility enables researchers to seamlessly tran- sition between different CA within single, unified framework (Table We invite readers to ex- perience CAX's capabilities firsthand by accessing our examples interactive notebooks in Google Colab, conveniently linked README of the repository http:://ithhbb.commmacccctttooor
3.1 ARCHITECTURE AND DESIGN
CAX introduces unifying framework for all cellular automata types. This flexible architecture is built upon two key components: the perceive module and the update module, see Figure 2 To- gether, these modules define the local rule of the CA. each time step, this local rule is applied uniformly to all cells in the grid, generating the next global state of the system, as explained in Sec- tion 2.1 This modular approach not only provides a clear separation of concerns but also facilitates easy experimentation and extension of existing CA models.
Building on the controllable CA framework introduced in Section 2.2 our architecture generalizes the neural cellular automata approach to recurrent convolutional neural architectures, enabling seam- less integration of external inputs and control signals. By implementing both modules using standard machine learning components from Flax, CAX makes it straightforward to experiment with various neural architectures while maintaining the cellular automata paradigm from simple convolutional layers to sophisticated attention mechanisms.
4

---

Published as a conference paper at ICLR 2025
@nnx. jit
def step (self, state: State, input: Input None """Perform a single step of the CA.
None) State:
Args:
state: Current state. input: Optional input.
Returns:
Updated state.
unn perception self perceive (state) state = self. update (state, perception, input> return state
The architecture of CAX allows for easy composition of different perceive and update modules, enabling the creation of wide variety of cellular automata models. This modular design also facilitates experimentation with new types of cellular automata by allowing users to define custom perceive and update modules while leveraging the existing infrastructure provided by the library.
3.1.1 PERCEIVE MODULE
The perceive module in CAX is responsible for gathering information from the neighborhood of each cell. This information then used by the update module determine the cell's next state. CAX provides several perception mechanisms, including Convolutional Perception, Depthwise Convolu- tional Perception and Fast Fourier Transform Perception. The perceive modules are designed to be flexible and can be customized for different types of cellular automata.
3.1.2 UPDATE MODULE
The update module in CAX is responsible for determining the next state of each cell based its current state and the information gathered by the perceive module. CAX provides several update mechanisms, including MI.P Update, Residual Update and Neural Cellular Automata Update. Like the perceive modules, the update modules are designed to be flexible and can be customized for different cellular automata models.
3.2 FEATURES
3.2.
PERFORMANCE
Cellular Automata Simulation
Neural Cellular Automata Training
140
CAX
Training steps per secand (stepu)s)
CAX
Staps per sacond (log stegus)
CellpyLb
120
101
100
10a
10a
10'
Elementary Colluaar Automata
Neural Cellular Automata
Self-clasatfying
Figure 3: Performance benchmarks of CAX. Left: Simulation speed comparison between CAX and CellPyLib for classical cellular automata. CAX demonstrates a 1, 1,400x speed-up for Flementary Cellular Automata and 2,000x speed-up for Conway's Game of Life. Right: Training speed comparison between CAX and the official TensorFlow implementation for neural cellular automata experiments. CAX achieves 1.5x speed-up on the Self-classifying MNIST Digits task.
5

---

Published as a conference paper at ICLR 2025
CAX leverages JAX's powerful vectorization and scan capabilities to achieve remarkable speed im- provements over existing implementations. Our benchmarks, conducted on a single NVIDIA RTX A6000 GPU, demonstrate significant performance gains across various cellular automata models. For Elementary Cellular Automata, CAX achieves a 1,400x speed-up compared to CellPyLib. In simulations of Conway's Game of Life, 2,000x speed-up observed relative Cell
Furthermore, the domain of Neural Cellular Automata, specifically the Self-classifying MNIST Digits experiment, CAX demonstrates 1.5x speed-up over the official TensorFlow implementation. These performance improvements, illustrated in Figure [ are made possible by JAX's efficient vectorization and the use of its sean operation for iterative computations. The following code snippet exemplifies how CAX utilizes JAX's s scan function to optimize multiple CA steps:
def step (carry: tuple(CA Statel, input: Input (CA, Statel, Statel: ca, state = carry
None) tuple (tuple
state = ca.step (state, input)
return (ca, state), state if all_steps else None
(, state), states = nax, scan
step, in_axes= (nnx. Carry, input in axis)
Length=num_stegg ((self, state), input)
This optimized approach allows for rapid execution of complex CA simulations, opening pos- sibilities for large-scale experiments and real-time applications.
3.2.2 UTILITIES
CAX offers rich set of utility functions support various aspects of cellular automata research. A high-quality implementation of the sampling pool technique is provided. which is crucial for training stable growing neural cellular automata (Mordvintsev et al.) 22020). To facilitate the training of unsu- pervised neural cellular automata and enable generative modeling within the CA framework, CAX incorporates a variational autoencoder implementation. Additionally, the library provides utilities for handling image and emoji inputs, allowing for diverse and visually engaging CA experiments. These utilities are designed to streamline common tasks in CA research, allowing researchers to focus on their specific experiments rather than reimplementing standard components.
3.2.3 DOCUMENTATION AND EXAMPLES
CAX prioritizes user experience and ease of adoption through comprehensive documentation and ex- amples. The entire library thoroughly documented, with typed classes and functions accompanied by descriptive docstrings. This ensures users have access to detailed information about CAX's s func- tionality and promotes clear, type-safe code. To help users get started and showcase advanced usage, CAX offers collection of tutorial-style interactive Colab notebooks. These notebooks demonstrate various applications of the library and can be run directly in a web browser without any prior setup, making it easy for new users to explore CAX's capabilities.
For easy access and integration into existing projects, CAX can be installed directly via PyPI, al- lowing users to quickly incorporate it into their Python environments. The library maintains high standards of code quality, with extensive unit tests covering significant portion the codebase. Continuous Integration (CI) pipelines ensure that all code changes are thoroughly tested and linted before integration. These features collectively make CAX not just a powerful tool for cellular au- tomata research, but also accessible and user-friendly library suitable for both novice and experi- enced researchers in the field.
IMPLEMENTED CELLULAR AUTOMATA AND EXPERIMENTS
To showcase the versatility and capabilities of the library, we show that CAX supports a wide array of cellular automata, ranging from classical discrete models advanced continuous CAs and in-

---

Published as a conference paper at ICLR 2025
cluding neural implementations. In this section, provide an overview of these implementations, demonstrating the library's flexibility handling various dimensions and types (Table D
We begin with three classic models that highlight CAX's ability to support both discrete and tinuous systems across different dimensions. The Elementary CA, a foundational one-dimensional discrete model studied extensively by Wolfram (2002), demonstrates CAX's efficiency in handling simple discrete systems. Conway's ' Game of Life (Gardner 1970p. a well-known two-dimensional model, showcases CAX's capability in simulating complex emergent behaviors in discrete space. Lenia (Chan 2019), a continuous, multi cilli-ddimenionall model, illustrates CAX's flexibility in sup- porting more complex, continuous systems in arbitrary dimensions.
Furthermore, we have replicated four prominent NCA experiments that have gained significant at- tention the field. The Growing NCA (Mordvintsev et al., 2020) demonstrates CAX's s ability to handle complex growing patterns and showcases the implementation of the sampling pool technique, crucial for stable growth and regeneration. The Growing Conditional NCA (Sudhakaran al. 2022) utilizes CAX's Controllable CA capabilities, introduced in Section 2.2 allowing for targeted pat- tern generation. The Growing Unsupervised NCA (Palm et al. 2021) highlights CAX's versatility in incorporating advanced machine learning techniques, specifically the use of Variational Au- toencoder within the NCA framework. The celf-classi MNIST Digits (Randazzo et al. 2020) showcases CAX's capacity for self-organe organizign systems with global coordination via local interac- tions, contrasting with growth-based tasks.
These implementations only validate CAX's performance and flexibility but also serve valu- able resources for researchers looking to build upon or extend these models. We complement these implementations with three novel experiments, which will be detailed in the following section.
Table 1: Overview of Cellular Automata implemented in CAX
Cellular Automata
Reference
Type
Dimensions
Elementary Cellular Automata
Wolfram 2002
Discrete
ID
Conway's 's Game Life
Gardner 1970
Discrete
Lenia
Chan 2019
Continuous
ND
Growing Neural Cellular Automata
Mordvintsev et al.
2020
Neural
2D
Growing Conditional Neural Cellular Automata
Suchakaran
2022
Neural
2D
Growing Unsupervised Neural Cellular Automata
Palm
Neural
2D
Self-classifying MNIST Digits
Randazzo et al.
2020
Neural
Diffusing Neural Cellular Automata
Section 5.1
Neural
2D
celf-auteencodin MNIST Digits
Section 5.2
Neural
3D
ID-ARC Neural Cellular Automata
Section
Neural
ID
NOVEL NEURAL. CELLULAR AUTOMATA EXPERIMENTS
5.1 DIFFUSING NEURAL CELLULAR AUTOMATA
In this experiment, introduce novel training procedure for NCA, inspired by diffusion models. Traditionally, NCAs have predominantly relied on growth- based training paradigms. where the state initialized with single alive cell and trained grow towards target pattern (Sudhakaran al., 2022 Mordvintsev et 2020 Palm et al. (021). However, this approach often faces challenges in maintaining stability and achieving consistent results (Mordvintsev al. 2020). The conventional
Figure 4: Inspired by diffusion models, the NCA learns to denoise images over a fixed number of steps. The process evolves from pure noise (left) to target pattern (right).

---

Published as a conference paper at ICLR 2025
NCA training method typically employs "sample pool" strategy to address stability issues and en- courage the formation of attractors. This approach involves maintaining diverse pool of interme- diate states, sampling from this pool for training, and periodically updating it with newly generated states. By exposing the NCA to various intermediate configurations and consistently guiding them towards the target pattern, the sample pool method helps shape the system's dynamics, making the desired pattern a more robust attractor in the state space.
Our proposed diffusion-inspired approach offers several advantages over the traditional growing mechanism. First, unlike the growing mechanism, our diffusion-basel approach does not require sample pool, which simplifies the training process and reduces memory requirements, making more efficient and scalable. Second, our diffusion-inspired approach naturally guides the NCA to- wards more stable dynamics, effectively creating stronger attractor basin around the target pattern. In Figure 5 we compare the regeneration capabilities of growing NCAs with diffusing NCAs. We create an artificial damage by cutting the tail of the gecko and observe that diffusing NCA demon- strate emergent regenerating capabilities.
Growing NCA
Diffusing NCA
Figure 5:
Diffusing NCAs demonstrate emergent regenerating capabilities compared to growing NCAs that unstable if not trained explicitely to regenerate and recover from damage.
5.2 SELF-AUTOENCODING MNIST DIGITS
In this experiment, we draw inspiration from Randazzo et al. (2020) where an NCA is trained classify MNIST digits through local inter- actions. In their work, each cell (pixel) of an MNIST digit learns output the correct digit label through local communication with neigh- boring cells. The NCA demonstrates the abil- ity to reach global consensus digit classi- fication, maintain this classification over time, and adapt perturbations mutations of the digit shape. Their model showcases emergent behavior, where simple local rules lead plex global patterns, analogous to biological systems achieving anatomical homeostasis.
Figure The 3D NCA is initialized with an MNIST digit (left). The NCA learns recon- struct the digit on the opposite red face (right).
Building upon this concept, we propose novel experiment that could be termed "Self-autoencodingg MNIST Digits". In this setup, we utilize three-dimensional NCA initialized with an MNIST digit on one face, see Figure 6 The objective of the NCA is to learn a rule that will replicate the MNIST digit on its opposite face (red face). However, we introduce critical constraint: the middle of the NCA, there is mask where cells cannot be updated, effectively preventing direct communication between the two faces. Crucially, we allow for a single-cell wide hole in the center of this mask, creating minimal channel for information transfer.
To successfully replicate the MNIST digit on the opposite face, the NCA must develop a sophis- ticated rule set that accomplishes two key tasks. First, it must encode the MNIST image into compressed form that can pass through the single-cell hole. Second, must then decode this infor- mation the other side accurately reconstruct the original digit. A notable aspect of this result is that each cell the NCA performs an identical local update rule, contributing to the system's 's over- all emergent behavior. As shown in Figure 7 the NCA successfully reconstructs MNIST digits
8

---

Copy Published as a conference paper at ICLR 2025
the red face, demonstrating its ability encode, transmit, and decode complex visual information through a minimal channel. This experiment highlights the power NCAs in learning complex information processing tasks using simple, uniform rules, while demonstrating CAX's ability to support sophisticated 3-dimensional CA rules.
/
2 3
4
5
6
7
8
9
D
/
2
3
4
6
7
3
9
Figure 7: The top shows the original digits from the test while the bottom row displays the corresponding reconstructions on the red face of the NCA.
5.3 1D-ARC NEURAL CELLULAR AUTOMATA
In this experiment, train one-dimensional NCA on the ID-ARC dataset (Xu al. 2024). The ID-ARC dataset is a novel adaptation of the original Abstraction and Reasoning Corpus (Chollet, 2019 (ARC), designed
simplify and streamline research in artificial intelligence and language models. By reducing the dimensionality of input and output images to single row of pixels, ID- ARC maintains the core knowledge priors of ARC while significantly reducing task complexity. For example, the tasks in ID-ARC include "Static movement by 3 pixels", "Fill", and "Recolor by Size Comparison". For full description of the dataset, see the project page Our experiment focuses on training NCA to solve the ID-ARC tasks. Each input sample consists of single row of colored pixels and a corresponding target row. The NCA's objective to transform the input into the target through successive applications of its rule. We consider task successful if all pixels in the NCA's output match the target pixels after a predetermined fixed number of steps.
Move
Move
Move
Move Dynamic
Move 2
Fil
Padded
Hallow
Flip
Minrar
Denaise
Danaisa
Fill
Multicolor
by Odd Even
by Size
by Comparison
Scaling
Copy
Multcolar
Figure 8: ID-ARC NCA space-time diagrams for each task. The top row of pixels in each image is the input. Subsequent rows of pixels show the NCA's intermediate steps as it attempts to transform the input into the target. The bottom row of pixels represents the NCA's final output after a fixed number of steps, which is compared to the target for task completion.

---

Published as a conference paper at ICLR 2025
The primary goal of this experiment is for the NCA to learn a generalizable rule from the training set, enabling it to solve unseen ex- amples in the test sets. This challenge tests the NCA's ability infer abstract patterns and apply them to new situations, key aspect of human-like reasoning. Figure & illlustrates the NCA's "reasoning" on all ID-ARC tasks. The visualization shows the input the top, in- termediate steps, and final output of the NCA the bottom each image, and is called space-time diagram.
Table GPT-4 and NCA aceuracy in percentage on all tasks from the ID-ARC test set. The GPT-4 val- ues are direct-grid approach, directly taken from Xu et al. 2024).
Task Move
GPT-4
NCA
66
100
Move
26
100
Move
24
100
Move Dynamic
22
12
Move Towards
34
98
Fill
66
66
To evaluate the NCA's performance, we com- pare it GPT-4, state-of-the-art language model, on the ID-ARC test set. Table - presents the accuracy of the NCA and GPT- across 18 different task types. The GPT- 4 values direct-grid results, directly taken from Xu et al. (2024). Notably, the NCA out- performs GPT-4 on several tasks, particularly those involving movement, pattern copying, and denoising. Overall, the NCA achieves total accuracy of 60. compared to GPT- 4's 41.56%, as reported by Xu al. (2024).
Padded Fill
26
28
Hollow
56
98
Flip
70
28
Mirror
20
Denoise
36
100
Denoise Multicolor
60
58
Pattern Copy
36
100
Pattern Copy Multicolor Recolor by Odd Even
38
100
32
Recolor by Size Recolor by Size Comparison Scaling
28
20
88
88
These results demonstrate the potential of NCAs solving abstract reasoning tasks, even outperforming sophisticated language models in certain domains. The NCA's success in tasks like "Move 3" and "Pattern Copy Multicolor" showcases its ability to learn complex spatial trans- formations and apply them consistently.
Total
41.56
60.12
However, the NCA struggles with tasks involving more abstract concepts like odd-even distinctions size comparisons. This limitation suggests areas for future improvement, possibly through the integration of additional priors or more sophisticated architectures. While the average of NCA outperforms GPT4, it interesting note that GPT4 performs equally in every task, while NCA completely fails on some of them (0% accuracy). This opens interesting questions for future work. This experiment not only highlights the capabilities of NCAs in abstract reasoning tasks but also demonstrates CAX's flexibility in implementing and training NCA models for diverse applications.
CONCLUSION
In this paper, we introduce CAX: Cellular Automata Accelerated in JAX, an open-source library, designed to provide high-performance and flexible framework accelerate cellular automata re- search. CAX provides substantial speed improvements over existing implementations, enabling researchers to run complex simulations and experiments more efficiently.
CAX's flexible architecture supports wide range of cellular automata types across multiple di- mensions, from classic discrete models advanced continuous and neural variants. Its modular design, based on customizable perceive and update components, facilitates rapid experimentation and development novel CA models, enabling efficient exploration of new ideas.
CAX's comprehensive documentation, example notebooks, and seamless integration with machine learning workflows not only lower the barrier to entry but also promote reproducibility and collabo- ration in cellular automata research. We hope this accessibility will accelerate the pace of discovery by attracting new researchers.
In the future, envision several exciting directions, such expanding the model zoo to implement and optimize a wider range cellular automata models, and exploring synergies between cellular automata and other approaches, such as reinforcement learning evolutionary algorithms.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Gabriel Béna and Arthur Braida for insightful discussions and feedback. We also ac- knowledge the financial support provided by grant from G-Research.
REFERENCES
Luis M. Antunes. Cellpylib: A python library for working with cellular automata. Journal of Open Source Software, 6(67);3608, 2021. doi: 10.21105/joss.0366 03600.. /RRLhttps:///i.ccrg
21105/ 03608
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miine and Qiao
Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. com ax ml/ ml/jax
Felix Chalumeau, Bryan Lim, Raphael Boige, Maxime Allard, Luca Grillotti, Manon Flageat, Valentin Macé, Arthur Flajolet, Thomas Pierrot, and Antoine Cully. QDax A Library for Quality- Diversity and Population- based Algorithms with Hardware Acceleration, August 2023. URL.
http://arxivv org/ abs 2308 03665 HEXiv:2308.0366: [es].
Bert Wang-Chak Chan. Lenia Biology of Artificial Life. Complex Systems, 28(3):251-288 October 2019. ISSN 08912513.
doi: 10.25000s/CmmpleeSyyttemm 28..
URL http:
/bspr////bb///// 1812 05433 1rXiv:1812.05433 [nlin].
Bert Wang-Chak Chan. Lenia and Expanded Universe. In The 2020 on pife.pp pp. 221-229, 2020. doi: 10. 10.1162/isal a 00297. htt.:ttpp//aax...or//a/////0//
arXiv:2005.03742 [nlin].
Bert Wang-Chak Chan. Towards Large-Scale Simulations of Open-Ended Evolution in Con- tinuous Cellular Automata, April 2023. URL http://arxxv.rrg/aa/330.....0
arXiv:2304.06639 [nlin].
François Chollet. On the Measure of Intelligence, November 2019. URL
abs/1911. 01547 1rXiv:1911.01547 [cs].
Matthew Cook. Universality in Elementary Cellular Automata. Complex Systems, 15(1):1-40,
March 2004. ISSN 08912513. doi: 10.5008SCCompeeeyttess...... URL htttps://wi wwww.
comp systems.t......... com wttpttaactt//lls1111.eee
DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Clau- dio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, lurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojeiech Stokowiee, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL htlp:
//github. com deepmind
Maxence Faldor and Antoine Cully. Toward artificial open- open-ended evolution within lenia using quality-diversiie Artificial Life, 2024.
Martin Gardner. Mathematical games. Scientific American, 223(4):120-123,, 1970. ISSN 00368733, 19467087. URL (tRLh/tp://wwwt http:://wwwwwsstt..oo.oooo wwtp://ogo//tttae////////e..
William Gilpin. Cellular automata convolutional neural networks. Physical Review E, 100(3):
032402, September 2019. ISSN 2470-0045, 2470-0053.
doi: 10.1 htts/PhyyReeE..00.0333.
org haap://bs/18000 02942 arXiv:1809.0299: [cond-mat, physics:nlin, physics:physics].
Morgan Giraud. Leniax: efficient and differentiable lenia simulator, 2022. stchub.cco/morggnnng.. htw.:/aaniaauu//// raud/lent
URL http://

---

2022 Published as a conference paper at ICLR 2025
Daniele Grattarola, Lorenzo Livi, and Cesare Alippi. Learning Graph Cellular Automata. In Ad-
vances in Neural Information Processing Systems, volume 34, pp. 20983 20994. Curran Asso- ciates, Inc., 2021. URL /ttppp::/ppooeeeingg..s.. www.iiss...e.iip....... cc www.erriss.cpppppee/e/eee/ af87f7cdcca223cc4 c3 oat.:a222ccccccaoaoeeeoaaaaaaaae Abtt
wwtppp/pp//////aaaa... Abst aasttract.... htmll
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2024, URL
J.H. Holland. Emergence: From Chaos to Order. Popular science / Oxford University Press. Oxford University Press, 2000. ISBN 9780192862122. URL https://books. gwwgle. co. uk/
Robert Tjarko Lange. evosax: IAs-based Evolution Strategies, December 2022. URL http:
//axxiv..rg/aa///2211.. 04180 arXiv:2212.44880 [cs].
Alexander Mordvintsev, Ettore Randazzo, Eyvind Niklasson, and Michael Levin. Growing neural cellular automata. Distill, 2020. doi: 10.2391 113311//iiiill0000222
Elias Najarro, Shyam Sudhakaran, Claire Glanois, and Sebastian Risi. HyperNCA: Growing Devel-
opmental Networks with Neural Cellular Automata, April 2022. http::/ttpp/ii...tttt
abs/ 2204 11674 arXiv:2204. 1674 11677 [cs].
John Von Neumann and Arthur W. Burks. Theory of Self-Reproducing Automata. University of Illinois Press, USA, 1966.
Ehsan Pajouheshgar, Yitao Xu, Tong Zhang, and Sabine Süsstrunk. DyNCA: Real-Time Dynamic Texture Synthesis Using Neural Cellular Automata. In 2023 IEEE/CVF Conference on Com- puter Vision and Puttern Recognition (CVPR), pp. 20742-20751, Vancouver, BC, Canada, June 2023. IEEE. ISBN 9798350301298,
doi: 10.1 1AII19/CCVR3222222220000022 URL https: //ieeexplore. www.:reeeeeeeeeeeeeeeeeooo hrg.://o......eeeeoooocccee////0
Rasmus Berg Palm, Miguel González Duque Shyam Sudhakaran, and Sebastian Risi Variational Neural Cellular Automata. October 2021. URL httphtts//ppeeree..eeet net. htt//ooormmoo
7FF04CMBx
Ettore Randazzo and Alexander Mordvintsev. Biomaker CA: a Biome Maker project using Cellular Automata, July 2023. URI. [cs].
09320 arXiv:2307.0933
Ettore Randazzo, Alexander Mordvintsev, Eyvind Niklasson, Michael Levin, and Sam Grey- danus. Self-classifying mnist digits. Distill, 2020. doi: https://dissill.pb//000///e/
Ettore Randazzo, Alexander Mordvintsev, Eyvind Niklasson, and Michael Levin. Adversarial reprogramming of neural cellular automata. Distill, 2021. doi: 11223911//isiill.0002..00
Shyam Sudhakaran, Elias Najarro, and Sebastian Risi. Goal-Guided Neural Cellular Automata: Learning to Control Self-Organising Systems, April 2022. URL UtLLittp::/aaxiv.oorr/////
2205. 06806 arXiv:2205. 006006 [cs].
Yujin Tang, Yingtao Tian, and David Ha. EvoJAX: Hardware-Acceleratte Neuroevolution. In Proceedings of the Genetic and Evolutionary Computation Conference Companion, pp. 308-311, July 2022. doi: 10. 145/3520304435587770 URUnttpp//arxxv..og//aaaa arXiv:2202.05008 [cs].
05008
Mattic Tesfaldet, Derek Nowrouzezahrai, and Chris Pal. Attention- Neural Cellular
Automata. Advances in Neural Information Processing Systems, 35:8174-8186, ber 2022. URL
Decem-
ngs nww.ins.....ee.. httpppss.ccppeeee wwtppppppppppeeeeeeeeee .tw.redu.....nnnnnnoo..e.... fttps///al/eeaa//222222 Mash/361es 2atp://cce1222caa001116666666622 -aaelbbbe-bbraactt..ccc Centtraact-oooffeennee.ffee html
Stephen Wolfram. New Kind of Science. Wolfram Media, 2002. ISBN 1579550088, URL
https: //wp:////ww.... wol Framscience com
12

---

Published as a conference paper at ICLR 2025
N. Wulff and J A Hertz. Learning Cellular Automaton Dynamics with Neural Net- works. In Advances in Neural Information Processing Systems, volume 5. Morgan- Kaufmann, 1992. URL https: stEps::/prrcceed.ng...out d6c651 ddcd971ddddno 21w117133224444444231201
Yudong Xu, Wenhao Li, Pashootan Vaezipoor, Scott Sanner, and Elias B. Khalil. LLMs and the Abstraction and Reasoning Corpus: Successes, Failures, and the Importance of Object- based Representations, February 2024. arXiv:2305. 18354 [es].
URL http:// /ttp///axvvvv///a wtt.://aa/......a..obo 18354
13

---

Published as a conference paper at ICLR 2025
A HYPERPARAMETERS
This appendix provides detailed hyperparameters for the three novel neural cellular automata (NCA) experiments introduced in Section 5 These hyperparameters govern the architecture, training process, and simulation dynamics of each experiment. For further details on the experimental setup, refer the respective subsections Section or the notebooks on
Table 3: Hyperparameters for Diffusing Neural Cellular Automata (see Section 5.1
Parameter
Value
Spatial dimensions
(72, 72)
Channel size
64
Number of kernels
3
Hidden size
256
Cell dropout rate
0.5
Batch size
8
Number of steps
64
Learning rate
0.001
Table 4: Hyperparameters for Self-autoencoding MNIST Digits (sce Section 5.2
Parameter
Value
Spatial dimensions
(16, 16, 32)
Channel size
32
Number of kernels
4
Hidden size
256
Cell dropout rate
0.5
Batch size
Number of steps
96
Learning rate
0.001
Pool size
1,024
Table 5: Hyperparameters for ID-ARC Neural Cellular Automata (see Section 5.3
Parameter
Value
Spatial dimensions
(96)
Channel size
64
Number of kernels
Hidden size
256
Cell dropout rate
0.5
Batch size
Number of steps
64
Learning rate
0.001
14


Reviews:
[{"review_id": "hlnY7XTNfp", "reviewer": "Reviewer_5JEh", "summary": "The paper present a new library that can model several types of cellular automata (CAs), ranging from simply Wolfram-style automata and Conway's game of life to continuous and/or complex neural cellular automata. This library is based on JAX and can thus utilize modern computation hardware better, leading to significant performance increases in some cases as the authors show on a few common examples for CA applications.", "questions": "None.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The issue the library tackles is impactful. CAs play a huge role in various areas of research and having efficient implementations is, of course, key to experimenters. Dramatic improvements in performance have been known to enable a way broader research community in other fields.\n\nThe library also unifies various scenarios in which CAs are used that have been fragmented over multiple platforms.\n\nThe foundation in JAX also seems to allow for comparatively easy integration with other approaches in the field of AI. (A discussion of that point might be helpful.)\n\nThe libraries capabilities are demonstrated on a well-selected range of applications.\n\nVarious interesting novel applications are defined and well presented in the paper.", "weaknesses": "The paper stresses some standard library features (like documentation or examples) too much compared to more research-relevant parts.\n\nThe performance comparison should be broader. Various libraries exist. Why were only two tested in only very specific cases?\n\nThe performance difference to TensorFlow is left unexplained. TF should use hardware acceleration as well, so what is happening there?\n\nThe figures are not properly referenced in the text.\n\nSome formal problems exist:\n- superfluous \")\" in \"Section 2)\" on page 4\n- ill-formatted citation of \"Mordvintsev et al.\" on page 6\n- many wrong quotation marks on pages 8-10\n- \"doesn't\" on page 8\n- \"a NCA\" on page 8\n- superfluous comma \"training set,\" page 10\n- ill-formatted reference to Appendix A on page 10", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "clu7cDkYWR", "reviewer": "Reviewer_BH4L", "summary": "This paper presents a software library for cellular automata with GPU acceleration through the JAX library. It demonstrates this new library on MNIST digit classification and the 1D version of the Abstract Reasoning Challenge", "questions": "Are the mentioned works feasible or simple to implement in CAX? Specifically ViTCA is interesting and could lead to an application of NCA on the full ARC set, but is the attention mechanism possible to implement in CAX?\n\nHow does the damage repair example of Figure 5 incorporate diffusion? Do you renoise the \"damaged\" image or is it simply through training on the diffusion process that it is able to recreate the original image? Could this process be used to train a generative NCA over a distribution of images instead of a single one?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The article is clearly written and makes a compelling argument for cellular automata. It demonstrates their utility in image analysis and even reasoning tasks. The use of GPU acceleration through JAX is timely and relevant as a new exploration of GPU acceleration capabilities. As the main contribution is a software library, it provides some experimental comparisons with other open-source libraries.\n\nThe secondary contribution of the article are three novel applications of NCA. The most concrete of these is the application to the 1D ARC dataset. This demonstrates that NCA can follow patterns from a small number of examples, allowing them to rival GPT4's in-context learning on visual reasoning tasks. To my knowledge, this is the first study of NCA for reasoning; the demonstration that NCA can do low-data pattern matching, whether or not that is \"reasoning\", is very interesting.", "weaknesses": "A question I had in reviewing this was its relevancy to ICLR. It isn't the standard ICLR paper, so I consulted ICLR's official CFP, which includes:\n\n+ generative models\n+ causal reasoning\n+ infrastructure, software libraries, hardware, etc.\n\nI believe the article fits into these categories, however the match to ICLR and the general machine learning literature could be strengthened. For example, these works use NCA in more traditional machine learning settings:\n\nGrattarola, Daniele, Lorenzo Livi, and Cesare Alippi. \"Learning graph cellular automata.\" Advances in Neural Information Processing Systems 34 (2021): 20983-20994.\n\nTesfaldet, Mattie, Derek Nowrouzezahrai, and Chris Pal. \"Attention-based neural cellular automata.\" Advances in Neural Information Processing Systems 35 (2022): 8174-8186.\n\nPajouheshgar, Ehsan, et al. \"Dynca: Real-time dynamic texture synthesis using neural cellular automata.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2023.\n\nNone of these are cited, despite being high-profile studies of NCA in machine learning (NeurIPS, CVPR). As the main contribution of this article is a software library, it would be very valuable to know if these studies or similar ones could be reproduced using CAX. For submission to ICLR, this article could do more to relate the contribution to machine learning.\n\nThe second major weakness is the limited experimental study. The studies proposed in section 5 function more as demonstrations of the CAX library than new experiments due to their lack of rigor. Only the 1D-ARC experiment presents quantified results and a comparison with another method. The novel diffusion method presented in 5.1 is only demonstrated on a single example, and the self-encoding MNIST digits experiment in 5.2 is only validated through visual comparison of 1 image per digit. As such, claims like \" As shown in Figure 7, the NCA successfully reconstructs MNIST digits on the red face, demonstrating its ability to encode, transmit, and decode complex visual information\nthrough a minimal channel\" are not fully supported. Experimental details and full reporting of experimental results are necessary to support such claims. As demonstrations of the CAX library, these examples are interesting, but as experiments, they fall short.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R41CXJ4NSH", "reviewer": "Reviewer_7WC1", "summary": "The paper presents a new open-source library for performing Cellular Automata studies written in the JAX framework. The use of JAX allows for significant acceleration of such studies, and the framework reduces the amount of boilerplate code scientists need to write when starting such an investigation. \n\nSeveral benchmarks are shown to demonstrate the structure of the software API and even demonstrate a result when compared with common language models.", "questions": "* Why not run proper hard and soft scaling tests in the performance comparison? It is always hard to show few numbers in a speed comparison. \n* Further, is the performance improvement only due to JAX or have you also written the code in a way that further improves parts of the algorithms? If so, that would make a key result in the paper and improve its novelty.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is very well written and easy to read. The authors have made presentation a clear priority in the paper and should be commended. \n\nThe software itself is also a great idea. Reducing the time it takes for people to start working on these kinds of studies and to do so in an accelerated, modern way is very useful and this work presents a framework doing exactly that.", "weaknesses": "While the paper presents a valuable tool, it does so in a very surface-level way. Most software discussions are in the API demonstration, which shows users how to put it into practice. This is important, but in an academic paper presenting the framework, the framework should take centre stage. Further, benchmarks against common test systems are very welcome in a software paper, but the comparison with GPT-4, while certainly interesting, is more of a research result in and of itself and not directly related to the software. I am sure there is more to the code that replacing numpy with JAX, but that should be laid out and explained in the manuscript. \n\nMy biggest concern is that ICLR isn't quite the correct platform for such a paper. I understand the desire for it to go there as it is related to ML but it is, at the end of the day, software. If the paper were to present a novel approach to structuring CA algorithms and packaging that up, it would be a different story. However, this decision is likely up to conference organizers. I can see the benefit of using ICLR as a platform to spread this software and perhaps help a lot of scientists. But from a newness perspective, it is not a breakthrough.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5wusf1r1Mw", "reviewer": "Reviewer_XZsB", "summary": "The authors have created an open-sourced python library called CAX (cellular automata accelerated in JAX) that implements a number of recent developments/experiments in the field of cellular automata as JAX models which allows for highly efficient and parallel computing. The authors also implement tools to customize these implementation for future researchers to run their own experiments in order to have a unified library to be used for efficient CA computations. Furthermore, they implement 3 new experiments with this library to demonstrate these customization tools (a diffusion CA model, an auto-encoding CA model, and neural CA model trained on 1D-ARC tasks).", "questions": "I don't have any major questions regarding this package other than the reason for the import crashes in the notebook I listed in the weaknesses.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Originality: While this paper does not necessarliy re-invent new methods, they have created an original and unified tool to be used by practitioners of cellular automata in a highly efficient way for a broad set of possible implementations. This is a tool that will surely accelerate the study of emergence and self-organization in cellular automata.\n\nQuality/Clarity: The paper is clearly written with sufficient background context to motivate the necessity of this tool as well as citing the relevant literature that has sparked recent interest in deeper neural cellular automata architectures.\n\nSignificance: Related to the originality of this tool, it is significant for practitioners of this field and lowers the barrier of entry for researchers to produce highly efficient code and run larger scale simulations and experiments. As the field of cellular automata is deeply intertwined with the study of emergence, scaling such simulations is a crucial requirement to push the field to new grounds. With such a tool, such experiments become much more feasible and cheaper for researchers.", "weaknesses": "- the performance increases plotted in figure 3 I feel are not sufficiently demonstrating the utility of such a library. I say this as a user of JAX and PyTorch and know that there are many performance increases when using JAX, and so there needs to be a much more comprehensive case being made here that JAX is actually introducing large speed ups. This is also increasingly important as there doesn't seem to be a massive speed up in the \"Growing cellular automata\" experiment, so perhaps the authors should compute performance changes for a larger set of experiments. (For example for the [self-organizing textures](https://distill.pub/selforg/2021/textures/) experiment, or [cellular automata controlling cartpole](https://arxiv.org/abs/2106.15240) experiment. \n\n- as someone that has had to write JAX code for my own custom CA implementations, it would be great if the appendix included some annotated code that explains how efficient JAX implementations can be written. I know that curious readers can always read the code on github, and there are some examples shown such as on line 281-289 as well as the example notebook in the appendix (but it has minimal comments), but it would be extremely helpful for new programmers to understand the thought process for writing such a library. This is doubly important because when trying to implement this library for custom experiments, a certain degree of programming in JAX will be required. This is possibly slightly outside the scope or responsibility of the authors, but I think it would go a long way for allowing researchers to adopt this library if there were tutorials or detailed annotations of code that explain this process (I know for myself this would have been very helpful a few years ago).\n\n- I very quickly tried to run one of the colab notebooks for the 1D-ARC tasks and had problems with importing. I did not test with all the notebooks but this should probably be double-checked across the board. I recieved two errors:\n   - (on the second cell: `%pip install \"cax[examples] @ git+https://github.com/879f4cf7/cax.git\"`): after installing most packages, I recieved the following message: `ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\ngcsfs 2024.10.0 requires fsspec==2024.10.0, but you have fsspec 2024.9.0 which is incompatible.`\n    - (on the third cell: `from cax.core.update.residual_update import ResidualUpdate`): \n~~~\n---------------------------------------------------------------------------\nModuleNotFoundError                       Traceback (most recent call last)\n<ipython-input-3-7bc79b5c20c3> in <cell line: 11>()\n      9 from cax.core.perceive.depthwise_conv_perceive import DepthwiseConvPerceive\n     10 from cax.core.perceive.kernels import grad_kernel, identity_kernel\n---> 11 from cax.core.update.residual_update import ResidualUpdate\n     12 from flax import nnx\n     13 from tqdm.auto import tqdm\n\n1 frames\n/usr/local/lib/python3.10/dist-packages/cax/core/update/mlp_update.py in <module>\n      5 import jax.numpy as jnp\n      6 from flax import nnx\n----> 7 from flax.nnx.nnx.nn import initializers\n      8 from flax.nnx.nnx.nn.linear import default_kernel_init\n      9 \n\nModuleNotFoundError: No module named 'flax.nnx.nnx'\n~~~", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)