A JAX implementation of Graph Rewriting Automata by Paul Cousin, inspired by this animation by Alexander Mordvintsev.

I wrote this partly to understand the GRA paper better (see my organic and handwritten notes), and partly because I wanted to replicate & play with the graphs like in Alex' animation - which I was not patient enough to wait for being published.

The implementation didn't end up faster than the original TensorFlow implementation. JAX's sparse support made parallelizing the division steps hard (dynamic shapes). In hindsight, adjacency matrices are the wrong representation ... neighbor lists make division O(1) per node. Buuut I managed to get the logic down to ~100 lines (grcax/rules.py).

My pre-opus 4.5 vibecoded animation attempt didn't make the cut, but luckily Alex' code is now public and it's awesome, go check it out!

Below are visualizations of the division step on simple graphs, which I used to try to make sense of how one could parallelize it:

Sequential divisions on K4 (complete graph, all nodes divided one by one):

Divisions on a 6-node checkerboard-adjacency graph (nodes 1, 4, 8):

The paper's main example rule. Each configuration (node state + neighbor states) maps to a new state and a division decision:

Evolution over 5 steps, starting from a minimal 3-regular graph: