Self-classifying MNIST Digits ¶

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 time

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

from cax.core import ComplexSystem
from cax.core.perceive import ConvPerceive
from cax.core.update import NCAUpdate
from cax.nn.pool import Pool
from cax.utils import clip_and_uint8

import time import jax import jax.numpy as jnp import mediapy import optax import torchvision from flax import nnx from jax import Array from cax.core import ComplexSystem from cax.core.perceive import ConvPerceive from cax.core.update import NCAUpdate from cax.nn.pool import Pool from cax.utils import clip_and_uint8

Configuration¶

In [4]:

seed = 0

spatial_dims = (28, 28)
channel_size = 20
perception_size = 80
hidden_layers_sizes = (80,)
cell_dropout_rate = 0.5

num_steps = 20
pool_size = 1_024
batch_size = 16
learning_rate = 1e-3

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

seed = 0 spatial_dims = (28, 28) channel_size = 20 perception_size = 80 hidden_layers_sizes = (80,) cell_dropout_rate = 0.5 num_steps = 20 pool_size = 1_024 batch_size = 16 learning_rate = 1e-3 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Dataset¶

In [5]:

# Load MNIST dataset
ds_train = torchvision.datasets.MNIST(root="./data", train=True, download=True)
ds_test = torchvision.datasets.MNIST(root="./data", train=False, download=True)

# Convert to jax.Array
x_train = jnp.array([x.resize(spatial_dims) for x, _ in ds_train])[..., None] / 255
x_test = jnp.array([x.resize(spatial_dims) for x, _ in ds_test])[..., None] / 255

y_integer_train = jnp.array([y for _, y in ds_train], dtype=jnp.int32)
y_integer_test = jnp.array([y for _, y in ds_test], dtype=jnp.int32)

# Visualize
mediapy.show_images(x_train[:8].repeat(4, axis=-3).repeat(4, axis=-2))

# Load MNIST dataset ds_train = torchvision.datasets.MNIST(root="./data", train=True, download=True) ds_test = torchvision.datasets.MNIST(root="./data", train=False, download=True) # Convert to jax.Array x_train = jnp.array([x.resize(spatial_dims) for x, _ in ds_train])[..., None] / 255 x_test = jnp.array([x.resize(spatial_dims) for x, _ in ds_test])[..., None] / 255 y_integer_train = jnp.array([y for _, y in ds_train], dtype=jnp.int32) y_integer_test = jnp.array([y for _, y in ds_test], dtype=jnp.int32) # Visualize mediapy.show_images(x_train[:8].repeat(4, axis=-3).repeat(4, axis=-2))

In [6]:

# fmt: off
color_lookup = jnp.array(
    [
        [128, 0, 0],      # Digit 0
        [230, 25, 75],    # Digit 1
        [70, 240, 240],   # Digit 2
        [210, 245, 60],   # Digit 3
        [250, 190, 190],  # Digit 4
        [170, 110, 40],   # Digit 5
        [170, 255, 195],  # Digit 6
        [165, 163, 159],  # Digit 7
        [0, 128, 128],    # Digit 8
        [128, 128, 0],    # Digit 9
        [0, 0, 0],        # Default
        [255, 255, 255],  # Background
    ]
) / 255


def compute_y(x, y_integer):
    """Compute the target y from image and integer label."""
    mask = x >= 0.1
    return jnp.where(mask, jax.nn.one_hot(y_integer, 10), 0.0)


def render(x, y):
    """Render x and y to RGB."""
    # Mask for digit and background pixels
    is_digit = (x > 0.1).astype(jnp.float32)
    is_not_digit = 1.0 - is_digit

    # Apply the mask to the probabilities
    y = is_digit * y

    black_and_white = jnp.concatenate([is_digit, is_not_digit], axis=-1) * 0.01
    y = jnp.concatenate([y, black_and_white], axis=-1)

    return color_lookup[jnp.argmax(y, axis=-1)]

# fmt: off color_lookup = jnp.array( [ [128, 0, 0], # Digit 0 [230, 25, 75], # Digit 1 [70, 240, 240], # Digit 2 [210, 245, 60], # Digit 3 [250, 190, 190], # Digit 4 [170, 110, 40], # Digit 5 [170, 255, 195], # Digit 6 [165, 163, 159], # Digit 7 [0, 128, 128], # Digit 8 [128, 128, 0], # Digit 9 [0, 0, 0], # Default [255, 255, 255], # Background ] ) / 255 def compute_y(x, y_integer): """Compute the target y from image and integer label.""" mask = x >= 0.1 return jnp.where(mask, jax.nn.one_hot(y_integer, 10), 0.0) def render(x, y): """Render x and y to RGB.""" # Mask for digit and background pixels is_digit = (x > 0.1).astype(jnp.float32) is_not_digit = 1.0 - is_digit # Apply the mask to the probabilities y = is_digit * y black_and_white = jnp.concatenate([is_digit, is_not_digit], axis=-1) * 0.01 y = jnp.concatenate([y, black_and_white], axis=-1) return color_lookup[jnp.argmax(y, axis=-1)]

In [7]:

y_train = jax.vmap(compute_y)(x_train, y_integer_train)
y_test = jax.vmap(compute_y)(x_test, y_integer_test)

y_train = jax.vmap(compute_y)(x_train, y_integer_train) y_test = jax.vmap(compute_y)(x_test, y_integer_test)

In [8]:

# Visualize different colored digits
digits = []
for i in range(10):
    mask = y_integer_train == i
    idx = jnp.argmax(mask)
    digits.append(render(x_train[idx], y_train[idx]))

mediapy.show_images(jnp.stack(digits).repeat(2, axis=-3).repeat(2, axis=-2))

# Visualize different colored digits digits = [] for i in range(10): mask = y_integer_train == i idx = jnp.argmax(mask) digits.append(render(x_train[idx], y_train[idx])) mediapy.show_images(jnp.stack(digits).repeat(2, axis=-3).repeat(2, axis=-2))

Instantiate system¶

In [9]:

class SelfClassifyingNCA(ComplexSystem):
    """Self-Classifying Neural Cellular Automata."""

    def __init__(self, *, rngs: nnx.Rngs):
        """Initialize Self-Classifying NCA.

        Args:
                rngs: rng key.

        """
        self.perceive = ConvPerceive(
            channel_size=channel_size,
            perception_size=perception_size,
            use_bias=True,
            activation_fn=nnx.relu,
            rngs=rngs,
        )
        self.update = NCAUpdate(
            channel_size=channel_size,
            perception_size=perception_size,
            hidden_layer_sizes=hidden_layers_sizes,
            cell_dropout_rate=cell_dropout_rate,
            zeros_init=True,
            rngs=rngs,
        )

    def _step(self, state: Array, input: Array | None = None) -> Array:
        """Perform a single step."""
        # Extract x
        x = state[..., -1:]

        # Step
        perception = self.perceive(state)
        next_state = self.update(state, perception, input)

        # Override
        next_state = next_state.at[..., -1:].set(x)

        return next_state

    @nnx.jit
    def render(self, state: Array):
        """Render state to RGB frame."""
        # Extract x and classification logits
        x = state[..., -1:]
        logits = state[..., :10]

        # Render the image and the logits to RGB
        rgb = render(x, logits)

        # Clip values to valid range and convert to uint8
        return clip_and_uint8(rgb)

class SelfClassifyingNCA(ComplexSystem): """Self-Classifying Neural Cellular Automata.""" def __init__(self, *, rngs: nnx.Rngs): """Initialize Self-Classifying NCA. Args: rngs: rng key. """ self.perceive = ConvPerceive( channel_size=channel_size, perception_size=perception_size, use_bias=True, activation_fn=nnx.relu, rngs=rngs, ) self.update = NCAUpdate( channel_size=channel_size, perception_size=perception_size, hidden_layer_sizes=hidden_layers_sizes, cell_dropout_rate=cell_dropout_rate, zeros_init=True, rngs=rngs, ) def _step(self, state: Array, input: Array | None = None) -> Array: """Perform a single step.""" # Extract x x = state[..., -1:] # Step perception = self.perceive(state) next_state = self.update(state, perception, input) # Override next_state = next_state.at[..., -1:].set(x) return next_state @nnx.jit def render(self, state: Array): """Render state to RGB frame.""" # Extract x and classification logits x = state[..., -1:] logits = state[..., :10] # Render the image and the logits to RGB rgb = render(x, logits) # Clip values to valid range and convert to uint8 return clip_and_uint8(rgb)

In [10]:

cs = SelfClassifyingNCA(rngs=rngs)

cs = SelfClassifyingNCA(rngs=rngs)

In [11]:

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

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

Number of params: 22580

Sample initial state¶

In [12]:

def sample_state(key):
    """Sample a state with a random image."""
    # Init state
    state = jnp.zeros(x_train.shape[1:3] + (channel_size,))

    # Sample random image
    x_idx = jax.random.choice(key, x_train.shape[0])
    x = x_train[x_idx]

    # Set image in state
    state = state.at[..., -1:].set(x)
    return state, x_idx

def sample_state(key): """Sample a state with a random image.""" # Init state state = jnp.zeros(x_train.shape[1:3] + (channel_size,)) # Sample random image x_idx = jax.random.choice(key, x_train.shape[0]) x = x_train[x_idx] # Set image in state state = state.at[..., -1:].set(x) return state, x_idx

Train¶

Pool¶

In [13]:

key, subkey = jax.random.split(key)

keys = jax.random.split(subkey, pool_size)
state, x_idx = jax.vmap(sample_state)(keys)

pool = Pool.create({"state": state, "x_idx": x_idx})

key, subkey = jax.random.split(key) keys = jax.random.split(subkey, pool_size) state, x_idx = jax.vmap(sample_state)(keys) pool = Pool.create({"state": state, "x_idx": x_idx})

Optimizer¶

In [14]:

lr_sched = optax.linear_schedule(
    init_value=learning_rate, end_value=0.01 * learning_rate, transition_steps=100_000
)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adam(learning_rate=lr_sched),
)

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

lr_sched = optax.linear_schedule( init_value=learning_rate, end_value=0.01 * learning_rate, transition_steps=100_000 ) optimizer = optax.chain( optax.clip_by_global_norm(1.0), optax.adam(learning_rate=lr_sched), ) optimizer = nnx.Optimizer(cs, optimizer, wrt=nnx.Param)

Loss¶

In [15]:

def l2(state, y):
    """L2."""
    l2_loss = jnp.sum(jnp.square(state[..., :10] - y), axis=(-1, -2, -3)) / 2
    return jnp.mean(l2_loss)


def ce(state, y):
    """Cross-entropy."""
    integer_label = jnp.argmax(y, axis=-1)
    return jnp.mean(
        optax.softmax_cross_entropy_with_integer_labels(state[..., :10], integer_label)
    )

def l2(state, y): """L2.""" l2_loss = jnp.sum(jnp.square(state[..., :10] - y), axis=(-1, -2, -3)) / 2 return jnp.mean(l2_loss) def ce(state, y): """Cross-entropy.""" integer_label = jnp.argmax(y, axis=-1) return jnp.mean( optax.softmax_cross_entropy_with_integer_labels(state[..., :10], integer_label) )

In [16]:

@nnx.jit
def loss_fn(cs, state, y):
    """Loss function."""
    state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
    state = nnx.split_rngs(splits=batch_size)(
        nnx.vmap(
            lambda cs, state: cs(state, num_steps=num_steps),
            in_axes=(state_axes, 0),
        )
    )(cs, state)

    loss = l2(state, y)
    return loss, state

@nnx.jit def loss_fn(cs, state, y): """Loss function.""" state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) state = nnx.split_rngs(splits=batch_size)( nnx.vmap( lambda cs, state: cs(state, num_steps=num_steps), in_axes=(state_axes, 0), ) )(cs, state) loss = l2(state, y) return loss, state

Train step¶

In [17]:

@nnx.jit
def train_step(cs, optimizer, pool, key):
    """Train step."""
    sample_key, sample_state_key = jax.random.split(key)

    # Sample from pool
    pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False)
    current_state = batch["state"]
    current_x_idx = batch["x_idx"]

    # A quarter of the batch is replaced with new images
    new_state, new_x_idx = sample_state(sample_state_key)
    current_state = current_state.at[: batch_size // 4].set(new_state)
    current_x_idx = current_x_idx.at[: batch_size // 4].set(new_x_idx)

    # Get images
    current_y = y_train[current_x_idx]

    (loss, current_state), grad = nnx.value_and_grad(loss_fn, has_aux=True)(
        cs, current_state, current_y
    )
    optimizer.update(cs, grad)

    pool = pool.update(pool_idx, {"state": current_state, "x_idx": current_x_idx})
    return loss, pool

@nnx.jit def train_step(cs, optimizer, pool, key): """Train step.""" sample_key, sample_state_key = jax.random.split(key) # Sample from pool pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False) current_state = batch["state"] current_x_idx = batch["x_idx"] # A quarter of the batch is replaced with new images new_state, new_x_idx = sample_state(sample_state_key) current_state = current_state.at[: batch_size // 4].set(new_state) current_x_idx = current_x_idx.at[: batch_size // 4].set(new_x_idx) # Get images current_y = y_train[current_x_idx] (loss, current_state), grad = nnx.value_and_grad(loss_fn, has_aux=True)( cs, current_state, current_y ) optimizer.update(cs, grad) pool = pool.update(pool_idx, {"state": current_state, "x_idx": current_x_idx}) return loss, pool

Main loop¶

In [18]:

num_train_steps = 8_192
print_interval = 128

losses = []
start = time.perf_counter()
for i in range(num_train_steps):
    key, subkey = jax.random.split(key)
    loss, pool = train_step(cs, optimizer, pool, subkey)
    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")

num_train_steps = 8_192 print_interval = 128 losses = [] start = time.perf_counter() for i in range(num_train_steps): key, subkey = jax.random.split(key) loss, pool = train_step(cs, optimizer, pool, subkey) 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")

/home/faldor_google_com/dev/cax/.venv/lib/python3.14/site-packages/jax/_src/interpreters/mlir.py:1286: UserWarning: A large amount of constants were captured during lowering (2.07GB total). If this is intentional, disable this warning by setting JAX_CAPTURED_CONSTANTS_WARN_BYTES=-1. To obtain a report of where these constants were encountered, set JAX_CAPTURED_CONSTANTS_REPORT_FRAMES=-1.
  warnings.warn(message)

W0904 21:46:59.569425  491972 cuda_timer.cc:88] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.

Step    0/8192 |   34.4s | Loss 6.428e+01

Step  128/8192 |   35.0s | Loss 6.514e+01

Step  256/8192 |   35.6s | Loss 6.558e+01

Step  384/8192 |   36.2s | Loss 6.420e+01

Step  512/8192 |   36.8s | Loss 6.273e+01

Step  640/8192 |   37.4s | Loss 6.318e+01

Step  768/8192 |   38.0s | Loss 6.065e+01

Step  896/8192 |   38.6s | Loss 6.118e+01

Step 1024/8192 |   39.2s | Loss 6.100e+01

Step 1152/8192 |   39.8s | Loss 6.044e+01

Step 1280/8192 |   40.4s | Loss 6.153e+01

Step 1408/8192 |   41.0s | Loss 6.158e+01

Step 1536/8192 |   41.6s | Loss 6.335e+01

Step 1664/8192 |   42.2s | Loss 6.334e+01

Step 1792/8192 |   42.8s | Loss 6.271e+01

Step 1920/8192 |   43.4s | Loss 6.320e+01

Step 2048/8192 |   44.0s | Loss 6.256e+01

Step 2176/8192 |   44.6s | Loss 6.333e+01

Step 2304/8192 |   45.2s | Loss 6.260e+01

Step 2432/8192 |   45.8s | Loss 6.232e+01

Step 2560/8192 |   46.4s | Loss 6.173e+01

Step 2688/8192 |   47.0s | Loss 6.152e+01

Step 2816/8192 |   47.6s | Loss 6.087e+01

Step 2944/8192 |   48.2s | Loss 6.011e+01

Step 3072/8192 |   48.8s | Loss 6.038e+01

Step 3200/8192 |   49.4s | Loss 6.031e+01

Step 3328/8192 |   50.0s | Loss 5.968e+01

Step 3456/8192 |   50.6s | Loss 5.785e+01

Step 3584/8192 |   51.2s | Loss 5.805e+01

Step 3712/8192 |   51.8s | Loss 5.672e+01

Step 3840/8192 |   52.4s | Loss 5.629e+01

Step 3968/8192 |   53.0s | Loss 5.491e+01

Step 4096/8192 |   53.6s | Loss 5.410e+01

Step 4224/8192 |   54.2s | Loss 5.344e+01

Step 4352/8192 |   54.8s | Loss 5.313e+01

Step 4480/8192 |   55.4s | Loss 5.156e+01

Step 4608/8192 |   56.0s | Loss 4.854e+01

Step 4736/8192 |   56.6s | Loss 4.822e+01

Step 4864/8192 |   57.2s | Loss 5.005e+01

Step 4992/8192 |   57.8s | Loss 4.853e+01

Step 5120/8192 |   58.4s | Loss 4.712e+01

Step 5248/8192 |   59.0s | Loss 4.678e+01

Step 5376/8192 |   59.6s | Loss 4.549e+01

Step 5504/8192 |   60.2s | Loss 4.447e+01

Step 5632/8192 |   60.8s | Loss 4.471e+01

Step 5760/8192 |   61.4s | Loss 4.628e+01

Step 5888/8192 |   62.0s | Loss 4.378e+01

Step 6016/8192 |   62.6s | Loss 4.128e+01

Step 6144/8192 |   63.2s | Loss 4.022e+01

Step 6272/8192 |   63.8s | Loss 4.163e+01

Step 6400/8192 |   64.5s | Loss 4.090e+01

Step 6528/8192 |   65.1s | Loss 3.958e+01

Step 6656/8192 |   65.7s | Loss 3.830e+01

Step 6784/8192 |   66.3s | Loss 3.956e+01

Step 6912/8192 |   66.9s | Loss 4.016e+01

Step 7040/8192 |   67.5s | Loss 3.728e+01

Step 7168/8192 |   68.1s | Loss 3.533e+01

Step 7296/8192 |   68.7s | Loss 3.659e+01

Step 7424/8192 |   69.3s | Loss 3.577e+01

Step 7552/8192 |   69.9s | Loss 3.444e+01

Step 7680/8192 |   70.5s | Loss 3.605e+01

Step 7808/8192 |   71.1s | Loss 3.609e+01

Step 7936/8192 |   71.7s | Loss 3.388e+01

Step 8064/8192 |   72.3s | Loss 3.714e+01

Step 8191/8192 |   72.9s | Loss 3.504e+01
✨ Trained for 8192 steps in 73s

Run¶

In [19]:

num_examples = 8

key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, num_examples)
state_init, _ = jax.vmap(sample_state)(keys)

state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
state_final, states = nnx.split_rngs(splits=num_examples)(
    nnx.vmap(
        lambda cs, state: cs(state, num_steps=4 * num_steps, return_states=True),
        in_axes=(state_axes, 0),
    )
)(cs, state_init)

num_examples = 8 key, subkey = jax.random.split(key) keys = jax.random.split(subkey, num_examples) state_init, _ = jax.vmap(sample_state)(keys) state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) state_final, states = nnx.split_rngs(splits=num_examples)( nnx.vmap( lambda cs, state: cs(state, num_steps=4 * num_steps, return_states=True), in_axes=(state_axes, 0), ) )(cs, state_init)

Visualize¶

In [20]:

states = jnp.concatenate([state_init[:, None], states], axis=1)
frames = nnx.vmap(
    lambda cs, state: cs.render(state),
    in_axes=(None, 0),
)(cs, states)

mediapy.show_videos(frames.repeat(4, axis=-3).repeat(4, axis=-2))

states = jnp.concatenate([state_init[:, None], states], axis=1) frames = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, states) mediapy.show_videos(frames.repeat(4, axis=-3).repeat(4, axis=-2))