Differentiable Logic Cellular Automata ¶

A neural cellular automaton learns a rule made of real numbers. This one learns a rule made of logic gates.

Each gate is one of the sixteen boolean functions of two inputs. That choice is discrete, so it cannot be differentiated. The trick is to relax it: during training a gate is a weighted average over all sixteen functions, with the weights a softmax over learned logits,

$$g(a, b) = \sum_{i=1}^{16} \operatorname{softmax}(w)_i \, f_i(a, b),$$

where each $f_i$ is written in a form that agrees with its boolean truth table on $\{0, 1\}$ and interpolates in between — $\mathrm{AND}(a,b) = ab$, $\mathrm{OR}(a,b) = a + b - ab$, $\mathrm{XOR}(a,b) = a + b - 2ab$, and so on.

At the end of training the softmax is replaced by an $\arg\max$. Every gate collapses to a single boolean function, and the automaton becomes an ordinary discrete circuit: no floating point, no approximation.

The target here is Conway's Game of Life. It is a rule we already know, which is the point — we can check the learned circuit against it exactly rather than judging it by eye.

Installation¶

You will need Python 3.12 or later, and a working JAX installation. For example, you can install JAX with:

In [1]:

%pip install -U "jax[cuda]"

%pip install -U "jax[cuda]"

/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip

Note: you may need to restart the kernel to use updated packages.

Then, install CAX from PyPi:

In [2]:

%pip install -U "cax[examples]"

%pip install -U "cax[examples]"

/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip

Note: you may need to restart the kernel to use updated packages.

Import¶

In [3]:

import itertools
import time

import jax
import jax.numpy as jnp
import mediapy
import optax
from flax import nnx
from jax import Array

from cax.core import ComplexSystem
from cax.core.perceive import Perceive
from cax.core.update import Update
from cax.utils import clip_and_uint8

import itertools import time import jax import jax.numpy as jnp import mediapy import optax from flax import nnx from jax import Array from cax.core import ComplexSystem from cax.core.perceive import Perceive from cax.core.update import Update from cax.utils import clip_and_uint8

Configuration¶

In [4]:

seed = 0

channel_size = 1
num_kernels = 16
perceive_layer_sizes = (9, 8, 4, 2, 1)
update_layer_sizes = (128,) * 16 + (128, 64, 32, 16, 8, 4, 2)

num_train_steps = 16_000
batch_size = 20
learning_rate = 0.05

num_gates = 16
pass_through_gate = 3
pass_through_logit = 10.0

key = jax.random.key(seed)
rngs = nnx.Rngs(seed)

seed = 0 channel_size = 1 num_kernels = 16 perceive_layer_sizes = (9, 8, 4, 2, 1) update_layer_sizes = (128,) * 16 + (128, 64, 32, 16, 8, 4, 2) num_train_steps = 16_000 batch_size = 20 learning_rate = 0.05 num_gates = 16 pass_through_gate = 3 pass_through_logit = 10.0 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Dataset¶

Conway's Game of Life, written out directly. A cell is born with exactly three live neighbours and survives with two or three; every other cell dies.

In [5]:

offsets = [(dy, dx) for dy in (-1, 0, 1) for dx in (-1, 0, 1)]


def life_step(board: Array) -> Array:
    """Apply one step of Conway's Game of Life on a periodic board."""
    neighbors = sum(
        jnp.roll(board, shift, (0, 1)) for shift in offsets if shift != (0, 0)
    )
    return ((neighbors == 3) | ((board == 1) & (neighbors == 2))).astype(board.dtype)

offsets = [(dy, dx) for dy in (-1, 0, 1) for dx in (-1, 0, 1)] def life_step(board: Array) -> Array: """Apply one step of Conway's Game of Life on a periodic board.""" neighbors = sum( jnp.roll(board, shift, (0, 1)) for shift in offsets if shift != (0, 0) ) return ((neighbors == 3) | ((board == 1) & (neighbors == 2))).astype(board.dtype)

A cellular automaton rule is a function of a $3 \times 3$ neighborhood, so there are only $2^9 = 512$ inputs it can ever see. Enumerating all of them gives a training set that covers the rule completely: a circuit that fits it has not generalized from examples, it has been shown every case.

On a $3 \times 3$ periodic board every cell sees the whole board, so these 512 boards present all 512 neighborhoods.

In [6]:

boards = (jnp.arange(512)[:, None] & (1 << jnp.arange(8, -1, -1))) > 0
boards = boards.reshape(512, 3, 3).astype(jnp.float32)

x = boards[..., None]
y = jax.vmap(life_step)(boards)[..., None]

mediapy.show_images([clip_and_uint8(b) for b in boards[:8]], width=64, height=64)

boards = (jnp.arange(512)[:, None] & (1 << jnp.arange(8, -1, -1))) > 0 boards = boards.reshape(512, 3, 3).astype(jnp.float32) x = boards[..., None] y = jax.vmap(life_step)(boards)[..., None] mediapy.show_images([clip_and_uint8(b) for b in boards[:8]], width=64, height=64)

Instantiate system¶

Gates¶

The sixteen boolean functions of two inputs, each written as a polynomial that reproduces its truth table on $\{0, 1\}$ and stays in $[0, 1]$ in between.

In [7]:

def binary_ops(a: Array, b: Array) -> Array:
    """Evaluate all sixteen two-input boolean functions relaxed to the unit interval."""
    return jnp.stack(
        [
            jnp.zeros_like(a),  # false
            a * b,  # a and b
            a - a * b,  # a and not b
            a,  # a
            b - a * b,  # b and not a
            b,  # b
            a + b - 2 * a * b,  # a xor b
            a + b - a * b,  # a or b
            1 - (a + b - a * b),  # a nor b
            1 - (a + b - 2 * a * b),  # a xnor b
            1 - b,  # not b
            1 - b + a * b,  # b implies a
            1 - a,  # not a
            1 - a + a * b,  # a implies b
            1 - a * b,  # a nand b
            jnp.ones_like(a),  # true
        ],
        axis=-1,
    )


def gate_layer(
    logits: Array, wire_a: Array, wire_b: Array, x: Array, *, hard: bool
) -> Array:
    """Apply one layer of learnable gates along the last axis of `x`.

    Under `hard`, each gate is the single boolean function its logits favor most, and
    the layer maps bits to bits. Otherwise it is the softmax-weighted average of all
    sixteen, which is differentiable and agrees with the hard version once the logits
    are decided.
    """
    weights = (
        jax.nn.one_hot(jnp.argmax(logits, axis=-1), num_gates)
        if hard
        else jax.nn.softmax(logits, axis=-1)
    )
    return jnp.sum(binary_ops(x[..., wire_a], x[..., wire_b]) * weights, axis=-1)

def binary_ops(a: Array, b: Array) -> Array: """Evaluate all sixteen two-input boolean functions relaxed to the unit interval.""" return jnp.stack( [ jnp.zeros_like(a), # false a * b, # a and b a - a * b, # a and not b a, # a b - a * b, # b and not a b, # b a + b - 2 * a * b, # a xor b a + b - a * b, # a or b 1 - (a + b - a * b), # a nor b 1 - (a + b - 2 * a * b), # a xnor b 1 - b, # not b 1 - b + a * b, # b implies a 1 - a, # not a 1 - a + a * b, # a implies b 1 - a * b, # a nand b jnp.ones_like(a), # true ], axis=-1, ) def gate_layer( logits: Array, wire_a: Array, wire_b: Array, x: Array, *, hard: bool ) -> Array: """Apply one layer of learnable gates along the last axis of `x`. Under `hard`, each gate is the single boolean function its logits favor most, and the layer maps bits to bits. Otherwise it is the softmax-weighted average of all sixteen, which is differentiable and agrees with the hard version once the logits are decided. """ weights = ( jax.nn.one_hot(jnp.argmax(logits, axis=-1), num_gates) if hard else jax.nn.softmax(logits, axis=-1) ) return jnp.sum(binary_ops(x[..., wire_a], x[..., wire_b]) * weights, axis=-1)

Wiring¶

Each gate reads two of the previous layer's outputs. Which two is fixed at initialization and never learned — only the choice of function is. The first perception layer is wired to the Moore neighborhood, pairing every neighbor with the center cell; deeper layers pair inputs so that no two gates in a layer see the same pair.

In [8]:

class Wire(nnx.Variable):
    """Fixed wiring between gate layers: carried with the model, never trained."""


def moore_wiring(key: Array) -> tuple[Array, Array]:
    """Pair every neighbor with the center cell."""
    neighbors = jnp.array([0, 1, 2, 3, 5, 6, 7, 8])
    perm = jax.random.permutation(key, neighbors.size)
    return neighbors[perm], jnp.full(neighbors.size, 4)[perm]


def unique_wiring(in_size: int, out_size: int, key: Array) -> tuple[Array, Array]:
    """Pair up inputs, widening the stride until there are enough distinct pairs."""
    index = jnp.arange(in_size)
    a, b = [index[::2]], [index[1::2]]
    pairs = min(a[0].size, b[0].size)
    a, b = [a[0][:pairs]], [b[0][:pairs]]
    stride = 1
    while sum(part.size for part in a) < out_size and stride < in_size:
        a.append(index[:-stride])
        b.append(index[stride:])
        stride += 1
    a, b = jnp.concatenate(a)[:out_size], jnp.concatenate(b)[:out_size]
    perm = jax.random.permutation(key, out_size)
    return a[perm], b[perm]


def init_logits(*shape: int) -> Array:
    """Start every gate as a pass-through, so an untrained circuit is the identity.

    The logit is large rather than merely largest: a circuit this deep, started from a
    near-uniform mixture of all sixteen functions, blurs its input into a constant
    before it reaches the output and leaves nothing to descend.
    """
    return (
        jnp.zeros(shape + (num_gates,))
        .at[..., pass_through_gate]
        .set(pass_through_logit)
    )

class Wire(nnx.Variable): """Fixed wiring between gate layers: carried with the model, never trained.""" def moore_wiring(key: Array) -> tuple[Array, Array]: """Pair every neighbor with the center cell.""" neighbors = jnp.array([0, 1, 2, 3, 5, 6, 7, 8]) perm = jax.random.permutation(key, neighbors.size) return neighbors[perm], jnp.full(neighbors.size, 4)[perm] def unique_wiring(in_size: int, out_size: int, key: Array) -> tuple[Array, Array]: """Pair up inputs, widening the stride until there are enough distinct pairs.""" index = jnp.arange(in_size) a, b = [index[::2]], [index[1::2]] pairs = min(a[0].size, b[0].size) a, b = [a[0][:pairs]], [b[0][:pairs]] stride = 1 while sum(part.size for part in a) < out_size and stride < in_size: a.append(index[:-stride]) b.append(index[stride:]) stride += 1 a, b = jnp.concatenate(a)[:out_size], jnp.concatenate(b)[:out_size] perm = jax.random.permutation(key, out_size) return a[perm], b[perm] def init_logits(*shape: int) -> Array: """Start every gate as a pass-through, so an untrained circuit is the identity. The logit is large rather than merely largest: a circuit this deep, started from a near-uniform mixture of all sixteen functions, blurs its input into a constant before it reaches the output and leaves nothing to descend. """ return ( jnp.zeros(shape + (num_gates,)) .at[..., pass_through_gate] .set(pass_through_logit) )

Perception and update¶

Perception runs num_kernels small circuits over the neighborhood in parallel, each reducing nine bits to one, and hands the update the cell's own state alongside their outputs. The update is a single deep circuit mapping that to the next state.

In [9]:

class LogicPerceive(Perceive[Array, Array]):
    """Read the Moore neighborhood through parallel logic circuits."""

    def __init__(
        self,
        *,
        channel_size: int,
        num_kernels: int,
        layer_sizes,
        hard: bool,
        rngs: nnx.Rngs,
    ):
        """Initialize the perception circuits.

        Args:
                channel_size: Number of state channels.
                num_kernels: Number of circuits read in parallel.
                layer_sizes: Width of each layer, starting at the nine neighborhood
                    cells.
                hard: Whether gates are discrete rather than softmax-weighted.
                rngs: rng key.

        """
        self.channel_size = channel_size
        self.num_kernels = num_kernels
        self.hard = hard
        keys = jax.random.split(rngs.params(), len(layer_sizes) - 1)
        self.logits, self.wires = nnx.List([]), nnx.List([])
        for i, (in_size, out_size) in enumerate(itertools.pairwise(layer_sizes)):
            wiring = (
                moore_wiring(keys[i])
                if i == 0
                else unique_wiring(in_size, out_size, keys[i])
            )
            self.logits.append(nnx.Param(init_logits(num_kernels, out_size)))
            self.wires.append(nnx.List([Wire(wiring[0]), Wire(wiring[1])]))

    def __call__(self, state: Array) -> Array:
        """Perceive the neighborhood of every cell.

        Args:
                state: State of the cellular automaton.

        Returns:
                The cell's own state followed by the output of each perception circuit.

        """
        neighborhood = jnp.stack(
            [jnp.roll(state, (-dy, -dx), (0, 1)) for dy, dx in offsets], axis=-1
        )
        x = jnp.broadcast_to(neighborhood, (self.num_kernels, *neighborhood.shape))
        for logits, wires in zip(self.logits, self.wires, strict=True):
            x = gate_layer(
                logits[:, None, None, None],
                wires[0][...],
                wires[1][...],
                x,
                hard=self.hard,
            )
        kernels = jnp.moveaxis(x, 0, -1).reshape(*state.shape[:2], -1)
        return jnp.concatenate([state, kernels], axis=-1)


class LogicUpdate(Update[Array, Array, Array]):
    """Map a perception to the next state through one deep logic circuit."""

    def __init__(self, *, layer_sizes, hard: bool, rngs: nnx.Rngs):
        """Initialize the update circuit.

        Args:
                layer_sizes: Width of each layer, from the perception to the state.
                hard: Whether gates are discrete rather than softmax-weighted.
                rngs: rng key.

        """
        self.hard = hard
        keys = jax.random.split(rngs.params(), len(layer_sizes) - 1)
        self.logits, self.wires = nnx.List([]), nnx.List([])
        for i, (in_size, out_size) in enumerate(itertools.pairwise(layer_sizes)):
            wire_a, wire_b = unique_wiring(in_size, out_size, keys[i])
            self.logits.append(nnx.Param(init_logits(out_size)))
            self.wires.append(nnx.List([Wire(wire_a), Wire(wire_b)]))

    def __call__(
        self, state: Array, perception: Array, input: Array | None = None
    ) -> Array:
        """Process the current state, perception, and input to produce a new state.

        Args:
                state: Current state.
                perception: Current perception.
                input: Optional input.

        Returns:
                Next state.

        """
        x = perception
        for logits, wires in zip(self.logits, self.wires, strict=True):
            x = gate_layer(logits, wires[0][...], wires[1][...], x, hard=self.hard)
        return x


class DiffLogicCA(ComplexSystem):
    """Differentiable Logic Cellular Automata class."""

    def __init__(self, *, hard: bool = False, rngs: nnx.Rngs):
        """Initialize Differentiable Logic CA.

        Args:
                hard: Whether gates are discrete rather than softmax-weighted.
                rngs: rng key.

        """
        self.perceive = LogicPerceive(
            channel_size=channel_size,
            num_kernels=num_kernels,
            layer_sizes=perceive_layer_sizes,
            hard=hard,
            rngs=rngs,
        )
        perception_size = channel_size * (1 + num_kernels * perceive_layer_sizes[-1])
        self.update = LogicUpdate(
            layer_sizes=(perception_size, *update_layer_sizes, channel_size),
            hard=hard,
            rngs=rngs,
        )

    def _step(self, state: Array, input: Array | None = None) -> Array:
        perception = self.perceive(state)
        next_state = self.update(state, perception, input)

        return next_state

    @nnx.jit
    def render(self, state):
        """Render state to RGB."""
        return clip_and_uint8(jnp.repeat(state, 3, axis=-1))

class LogicPerceive(Perceive[Array, Array]): """Read the Moore neighborhood through parallel logic circuits.""" def __init__( self, *, channel_size: int, num_kernels: int, layer_sizes, hard: bool, rngs: nnx.Rngs, ): """Initialize the perception circuits. Args: channel_size: Number of state channels. num_kernels: Number of circuits read in parallel. layer_sizes: Width of each layer, starting at the nine neighborhood cells. hard: Whether gates are discrete rather than softmax-weighted. rngs: rng key. """ self.channel_size = channel_size self.num_kernels = num_kernels self.hard = hard keys = jax.random.split(rngs.params(), len(layer_sizes) - 1) self.logits, self.wires = nnx.List([]), nnx.List([]) for i, (in_size, out_size) in enumerate(itertools.pairwise(layer_sizes)): wiring = ( moore_wiring(keys[i]) if i == 0 else unique_wiring(in_size, out_size, keys[i]) ) self.logits.append(nnx.Param(init_logits(num_kernels, out_size))) self.wires.append(nnx.List([Wire(wiring[0]), Wire(wiring[1])])) def __call__(self, state: Array) -> Array: """Perceive the neighborhood of every cell. Args: state: State of the cellular automaton. Returns: The cell's own state followed by the output of each perception circuit. """ neighborhood = jnp.stack( [jnp.roll(state, (-dy, -dx), (0, 1)) for dy, dx in offsets], axis=-1 ) x = jnp.broadcast_to(neighborhood, (self.num_kernels, *neighborhood.shape)) for logits, wires in zip(self.logits, self.wires, strict=True): x = gate_layer( logits[:, None, None, None], wires[0][...], wires[1][...], x, hard=self.hard, ) kernels = jnp.moveaxis(x, 0, -1).reshape(*state.shape[:2], -1) return jnp.concatenate([state, kernels], axis=-1) class LogicUpdate(Update[Array, Array, Array]): """Map a perception to the next state through one deep logic circuit.""" def __init__(self, *, layer_sizes, hard: bool, rngs: nnx.Rngs): """Initialize the update circuit. Args: layer_sizes: Width of each layer, from the perception to the state. hard: Whether gates are discrete rather than softmax-weighted. rngs: rng key. """ self.hard = hard keys = jax.random.split(rngs.params(), len(layer_sizes) - 1) self.logits, self.wires = nnx.List([]), nnx.List([]) for i, (in_size, out_size) in enumerate(itertools.pairwise(layer_sizes)): wire_a, wire_b = unique_wiring(in_size, out_size, keys[i]) self.logits.append(nnx.Param(init_logits(out_size))) self.wires.append(nnx.List([Wire(wire_a), Wire(wire_b)])) def __call__( self, state: Array, perception: Array, input: Array | None = None ) -> Array: """Process the current state, perception, and input to produce a new state. Args: state: Current state. perception: Current perception. input: Optional input. Returns: Next state. """ x = perception for logits, wires in zip(self.logits, self.wires, strict=True): x = gate_layer(logits, wires[0][...], wires[1][...], x, hard=self.hard) return x class DiffLogicCA(ComplexSystem): """Differentiable Logic Cellular Automata class.""" def __init__(self, *, hard: bool = False, rngs: nnx.Rngs): """Initialize Differentiable Logic CA. Args: hard: Whether gates are discrete rather than softmax-weighted. rngs: rng key. """ self.perceive = LogicPerceive( channel_size=channel_size, num_kernels=num_kernels, layer_sizes=perceive_layer_sizes, hard=hard, rngs=rngs, ) perception_size = channel_size * (1 + num_kernels * perceive_layer_sizes[-1]) self.update = LogicUpdate( layer_sizes=(perception_size, *update_layer_sizes, channel_size), hard=hard, rngs=rngs, ) def _step(self, state: Array, input: Array | None = None) -> Array: perception = self.perceive(state) next_state = self.update(state, perception, input) return next_state @nnx.jit def render(self, state): """Render state to RGB.""" return clip_and_uint8(jnp.repeat(state, 3, axis=-1))

In [10]:

cs = DiffLogicCA(rngs=rngs)

cs = DiffLogicCA(rngs=rngs)

In [11]:

params = nnx.state(cs, nnx.Param)
print("Number of gates:", sum(x.size // num_gates for x in jax.tree.leaves(params)))

params = nnx.state(cs, nnx.Param) print("Number of gates:", sum(x.size // num_gates for x in jax.tree.leaves(params)))

Number of gates: 2543

Train¶

Optimizer¶

In [12]:

optimizer = optax.chain(
    optax.clip(100.0),
    optax.adamw(learning_rate=learning_rate, b1=0.9, b2=0.99, weight_decay=1e-2),
)

optimizer = nnx.Optimizer(cs, optimizer, wrt=nnx.Param)

optimizer = optax.chain( optax.clip(100.0), optax.adamw(learning_rate=learning_rate, b1=0.9, b2=0.99, weight_decay=1e-2), ) optimizer = nnx.Optimizer(cs, optimizer, wrt=nnx.Param)

Loss¶

In [13]:

def loss_fn(cs, x, y):
    """Loss function."""
    prediction = nnx.vmap(lambda cs, state: cs(state, num_steps=1), in_axes=(None, 0))(
        cs, x
    )
    return jnp.sum(jnp.square(prediction - y))

def loss_fn(cs, x, y): """Loss function.""" prediction = nnx.vmap(lambda cs, state: cs(state, num_steps=1), in_axes=(None, 0))( cs, x ) return jnp.sum(jnp.square(prediction - y))

Train step¶

In [14]:

@nnx.jit
def train_step(cs, optimizer, x, y):
    """Train step."""
    loss, grad = nnx.value_and_grad(loss_fn)(cs, x, y)
    optimizer.update(cs, grad)

    return loss

@nnx.jit def train_step(cs, optimizer, x, y): """Train step.""" loss, grad = nnx.value_and_grad(loss_fn)(cs, x, y) optimizer.update(cs, grad) return loss

Main loop¶

In [15]:

print_interval = 800

losses = []
start = time.perf_counter()
for i in range(num_train_steps):
    key, subkey = jax.random.split(key)
    idx = jax.random.randint(subkey, (batch_size,), 0, x.shape[0])
    loss = train_step(cs, optimizer, x[idx], y[idx])
    losses.append(loss)

    if i % print_interval == 0 or i == num_train_steps - 1:
        avg_loss = sum(losses[-print_interval:]) / len(losses[-print_interval:])
        elapsed = time.perf_counter() - start
        print(f"Step {i:>4}/{num_train_steps} | {elapsed:6.1f}s | Loss {avg_loss:.3e}")

print(f"✨ Trained for {num_train_steps} steps in {time.perf_counter() - start:.0f}s")

print_interval = 800 losses = [] start = time.perf_counter() for i in range(num_train_steps): key, subkey = jax.random.split(key) idx = jax.random.randint(subkey, (batch_size,), 0, x.shape[0]) loss = train_step(cs, optimizer, x[idx], y[idx]) losses.append(loss) if i % print_interval == 0 or i == num_train_steps - 1: avg_loss = sum(losses[-print_interval:]) / len(losses[-print_interval:]) elapsed = time.perf_counter() - start print(f"Step {i:>4}/{num_train_steps} | {elapsed:6.1f}s | Loss {avg_loss:.3e}") print(f"✨ Trained for {num_train_steps} steps in {time.perf_counter() - start:.0f}s")

Step    0/16000 |   11.0s | Loss 9.213e+01

Step  800/16000 |   16.4s | Loss 2.677e+01

Step 1600/16000 |   21.7s | Loss 2.084e+00

Step 2400/16000 |   27.0s | Loss 1.333e-01

Step 3200/16000 |   32.4s | Loss 1.371e-02

Step 4000/16000 |   37.8s | Loss 4.125e-04

Step 4800/16000 |   43.2s | Loss 1.203e-06

Step 5600/16000 |   48.5s | Loss 5.071e-07

Step 6400/16000 |   53.9s | Loss 4.524e-07

Step 7200/16000 |   59.4s | Loss 4.495e-07

Step 8000/16000 |   64.7s | Loss 4.507e-07

Step 8800/16000 |   69.9s | Loss 4.488e-07

Step 9600/16000 |   75.2s | Loss 4.438e-07

Step 10400/16000 |   80.5s | Loss 4.456e-07

Step 11200/16000 |   85.8s | Loss 4.441e-07

Step 12000/16000 |   91.1s | Loss 4.432e-07

Step 12800/16000 |   96.4s | Loss 4.451e-07

Step 13600/16000 |  101.7s | Loss 4.439e-07

Step 14400/16000 |  107.0s | Loss 4.422e-07

Step 15200/16000 |  112.2s | Loss 4.437e-07

Step 15999/16000 |  117.4s | Loss 4.427e-07
✨ Trained for 16000 steps in 117s

Run¶

Training is over, so the softmax can go. Each gate becomes the single boolean function its logits favor, and the automaton stops being an approximation of a circuit and becomes one.

In [16]:

cs_hard = DiffLogicCA(hard=True, rngs=nnx.Rngs(seed))
nnx.update(cs_hard, nnx.state(cs))

prediction = nnx.vmap(lambda cs, state: cs(state, num_steps=1), in_axes=(None, 0))(
    cs_hard, x
)
correct = int(jnp.sum(jnp.abs(prediction - y) < 1e-6))
print(f"{correct}/{y.size} cells correct across all {x.shape[0]} neighborhoods")

cs_hard = DiffLogicCA(hard=True, rngs=nnx.Rngs(seed)) nnx.update(cs_hard, nnx.state(cs)) prediction = nnx.vmap(lambda cs, state: cs(state, num_steps=1), in_axes=(None, 0))( cs_hard, x ) correct = int(jnp.sum(jnp.abs(prediction - y) < 1e-6)) print(f"{correct}/{y.size} cells correct across all {x.shape[0]} neighborhoods")

4608/4608 cells correct across all 512 neighborhoods

Every neighborhood the rule can ever meet is in that set, so a perfect score is not an estimate of accuracy — it says the learned circuit is the Game of Life. It should then agree on a board far larger than anything it was trained on.

In [17]:

board = jax.random.randint(jax.random.key(1), (64, 64, 1), 0, 2).astype(jnp.float32)

learned, truth = board, board
frames = [cs_hard.render(board)]
for _ in range(64):
    learned = cs_hard(learned, num_steps=1)
    truth = life_step(truth[..., 0])[..., None]
    frames.append(cs_hard.render(learned))

print("identical to the Game of Life:", bool(jnp.array_equal(learned, truth)))

board = jax.random.randint(jax.random.key(1), (64, 64, 1), 0, 2).astype(jnp.float32) learned, truth = board, board frames = [cs_hard.render(board)] for _ in range(64): learned = cs_hard(learned, num_steps=1) truth = life_step(truth[..., 0])[..., None] frames.append(cs_hard.render(learned)) print("identical to the Game of Life:", bool(jnp.array_equal(learned, truth)))

identical to the Game of Life: True

Visualize¶

In [18]:

mediapy.show_video(jnp.stack(frames).repeat(4, axis=-3).repeat(4, axis=-2), fps=10)

mediapy.show_video(jnp.stack(frames).repeat(4, axis=-3).repeat(4, axis=-2), fps=10)