1D-ARC Neural Cellular Automata ¶

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 json
import time
from pathlib import Path

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 ConvPerceive, grad_kernel, identity_kernel
from cax.core.update import ResidualUpdate
from cax.utils import clip_and_uint8

import json import time from pathlib import Path 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 ConvPerceive, grad_kernel, identity_kernel from cax.core.update import ResidualUpdate from cax.utils import clip_and_uint8

Configuration¶

In [4]:

seed = 0

num_spatial_dims = 1
channel_size = 64
num_kernels = 2
hidden_layer_sizes = (256,)
cell_dropout_rate = 0.0

num_embeddings = 10  # 10 colors in total
features = 3  # embed in rgb

num_steps = 64
batch_size = 16
learning_rate = 1e-3

ds_size = 96

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

seed = 0 num_spatial_dims = 1 channel_size = 64 num_kernels = 2 hidden_layer_sizes = (256,) cell_dropout_rate = 0.0 num_embeddings = 10 # 10 colors in total features = 3 # embed in rgb num_steps = 64 batch_size = 16 learning_rate = 1e-3 ds_size = 96 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Dataset¶

In [5]:

!git clone https://github.com/khalil-research/1D-ARC.git

!git clone https://github.com/khalil-research/1D-ARC.git

fatal: destination path '1D-ARC' already exists and is not an empty directory.

/home/faldor_google_com/.local/share/uv/python/cpython-3.14.7-linux-x86_64-gnu/lib/python3.14/pty.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
  pid, fd = os.forkpty()

In [6]:

ds_path = Path("./1D-ARC/dataset")


def process(input, output):
    """Process input and output from dataset."""
    input = jnp.squeeze(jnp.array(input, dtype=jnp.int32))
    output = jnp.squeeze(jnp.array(output, dtype=jnp.int32))
    assert input.shape == output.shape

    pad_size = ds_size - input.size
    pad_left, pad_right = pad_size // 2, pad_size - pad_size // 2

    input_padded = jnp.pad(input, (pad_left, pad_right))
    output_padded = jnp.pad(output, (pad_left, pad_right))

    return jnp.stack([input_padded, output_padded])


ds = []
tasks = []
for task_idx, task_path in enumerate(ds_path.iterdir()):
    task_name = task_path.name
    for task_file in task_path.iterdir():
        with task_file.open() as f:
            data = json.load(f)
            input_output = jnp.array(
                [
                    process(data["train"][0]["input"], data["train"][0]["output"]),
                    process(data["train"][1]["input"], data["train"][1]["output"]),
                    process(data["train"][2]["input"], data["train"][2]["output"]),
                    process(data["test"][0]["input"], data["test"][0]["output"]),
                ],
                dtype=jnp.int32,
            )
            tasks.append(task_name)
            ds.append(input_output)
ds = jnp.stack(ds)

unique_tasks = list(set(tasks))
task_to_idx = {task: idx for idx, task in enumerate(unique_tasks)}
tasks = jnp.array([task_to_idx[task] for task in tasks], dtype=jnp.int32)

ds_path = Path("./1D-ARC/dataset") def process(input, output): """Process input and output from dataset.""" input = jnp.squeeze(jnp.array(input, dtype=jnp.int32)) output = jnp.squeeze(jnp.array(output, dtype=jnp.int32)) assert input.shape == output.shape pad_size = ds_size - input.size pad_left, pad_right = pad_size // 2, pad_size - pad_size // 2 input_padded = jnp.pad(input, (pad_left, pad_right)) output_padded = jnp.pad(output, (pad_left, pad_right)) return jnp.stack([input_padded, output_padded]) ds = [] tasks = [] for task_idx, task_path in enumerate(ds_path.iterdir()): task_name = task_path.name for task_file in task_path.iterdir(): with task_file.open() as f: data = json.load(f) input_output = jnp.array( [ process(data["train"][0]["input"], data["train"][0]["output"]), process(data["train"][1]["input"], data["train"][1]["output"]), process(data["train"][2]["input"], data["train"][2]["output"]), process(data["test"][0]["input"], data["test"][0]["output"]), ], dtype=jnp.int32, ) tasks.append(task_name) ds.append(input_output) ds = jnp.stack(ds) unique_tasks = list(set(tasks)) task_to_idx = {task: idx for idx, task in enumerate(unique_tasks)} tasks = jnp.array([task_to_idx[task] for task in tasks], dtype=jnp.int32)

In [7]:

import time

from PIL import ImageColor

# ARC-AGI colors
colors = {
    0: "#000000",  # Black
    1: "#0074D9",  # Blue
    2: "#FF4136",  # Red
    3: "#2ECC40",  # Green
    4: "#FFDC00",  # Yellow
    5: "#AAAAAA",  # Grey
    6: "#F012BE",  # Fuchsia
    7: "#FF851B",  # Orange
    8: "#7FDBFF",  # Teal
    9: "#870C25",  # Brown
}

# Convert all ARC colors to RGB using PIL
color_lookup = jnp.array([ImageColor.getrgb(hex) for hex in colors.values()]) / 255

import time from PIL import ImageColor # ARC-AGI colors colors = { 0: "#000000", # Black 1: "#0074D9", # Blue 2: "#FF4136", # Red 3: "#2ECC40", # Green 4: "#FFDC00", # Yellow 5: "#AAAAAA", # Grey 6: "#F012BE", # Fuchsia 7: "#FF851B", # Orange 8: "#7FDBFF", # Teal 9: "#870C25", # Brown } # Convert all ARC colors to RGB using PIL color_lookup = jnp.array([ImageColor.getrgb(hex) for hex in colors.values()]) / 255

In [8]:

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

tasks = jax.random.permutation(subkey, tasks)
ds = jax.random.permutation(subkey, ds)

split = int(0.9 * ds.shape[0])

train_ds = ds[:split]
train_tasks = tasks[:split]

test_ds = ds[split:]
test_tasks = tasks[split:]

key, subkey = jax.random.split(key) tasks = jax.random.permutation(subkey, tasks) ds = jax.random.permutation(subkey, ds) split = int(0.9 * ds.shape[0]) train_ds = ds[:split] train_tasks = tasks[:split] test_ds = ds[split:] test_tasks = tasks[split:]

Sample initial state¶

In [9]:

def create_state_with_sample(cs, sample, key):
    """Create state with sample."""
    # Sample input and target
    (
        (input_embed_1, output_embed_1),
        (input_embed_2, output_embed_2),
        (input_embed_3, output_embed_3),
        (input_embed, _),
    ) = cs.embed_input(sample)

    # Create context
    context_1 = jnp.concatenate(
        [
            input_embed_1,
            output_embed_1,
            input_embed_2,
            output_embed_2,
            input_embed_3,
            output_embed_3,
        ],
        axis=-1,
    )
    context_2 = jnp.concatenate(
        [
            input_embed_1,
            output_embed_1,
            input_embed_3,
            output_embed_3,
            input_embed_2,
            output_embed_2,
        ],
        axis=-1,
    )
    context_3 = jnp.concatenate(
        [
            input_embed_2,
            output_embed_2,
            input_embed_1,
            output_embed_1,
            input_embed_3,
            output_embed_3,
        ],
        axis=-1,
    )
    context_4 = jnp.concatenate(
        [
            input_embed_3,
            output_embed_3,
            input_embed_1,
            output_embed_1,
            input_embed_2,
            output_embed_2,
        ],
        axis=-1,
    )
    context_5 = jnp.concatenate(
        [
            input_embed_2,
            output_embed_2,
            input_embed_3,
            output_embed_3,
            input_embed_1,
            output_embed_1,
        ],
        axis=-1,
    )
    context_6 = jnp.concatenate(
        [
            input_embed_3,
            output_embed_3,
            input_embed_2,
            output_embed_2,
            input_embed_1,
            output_embed_1,
        ],
        axis=-1,
    )
    context = jax.random.choice(
        key,
        jnp.array([context_1, context_2, context_3, context_4, context_5, context_6]),
    )

    # Initialize state
    state = jnp.zeros((ds_size, channel_size))
    state = state.at[..., 3 : 18 + 3].set(context)
    state = state.at[..., -10:].set(jax.nn.one_hot(sample[-1, 0], num_classes=10))
    return state, sample[-1, -1]


def sample_state(cs, key):
    """Sample state with data augmentation."""
    key_sample, key_flip, key_perm, key_init = jax.random.split(key, 4)

    # Sample dataset
    _ = jax.random.choice(key_sample, train_tasks)
    sample = jax.random.choice(key_sample, train_ds)

    # Flip sample half of the time
    flip = jax.random.bernoulli(key_flip, p=0.5)
    sample = jnp.where(flip < 0.5, sample, jnp.flip(sample, axis=-1))

    # Permute colors
    color_perm = jnp.concatenate(
        [
            jnp.array([0], dtype=jnp.int32),
            jax.random.permutation(key_perm, jnp.arange(9)) + 1,
        ]
    )
    sample = color_perm[sample]

    return create_state_with_sample(cs, sample, key_init)


def sample_state_test(cs, key):
    """Sample state with data augmentation."""
    key_sample, key_init = jax.random.split(key)

    # Sample dataset
    _ = jax.random.choice(key_sample, test_tasks)
    sample = jax.random.choice(key_sample, test_ds)

    return create_state_with_sample(cs, sample, key_init)

def create_state_with_sample(cs, sample, key): """Create state with sample.""" # Sample input and target ( (input_embed_1, output_embed_1), (input_embed_2, output_embed_2), (input_embed_3, output_embed_3), (input_embed, _), ) = cs.embed_input(sample) # Create context context_1 = jnp.concatenate( [ input_embed_1, output_embed_1, input_embed_2, output_embed_2, input_embed_3, output_embed_3, ], axis=-1, ) context_2 = jnp.concatenate( [ input_embed_1, output_embed_1, input_embed_3, output_embed_3, input_embed_2, output_embed_2, ], axis=-1, ) context_3 = jnp.concatenate( [ input_embed_2, output_embed_2, input_embed_1, output_embed_1, input_embed_3, output_embed_3, ], axis=-1, ) context_4 = jnp.concatenate( [ input_embed_3, output_embed_3, input_embed_1, output_embed_1, input_embed_2, output_embed_2, ], axis=-1, ) context_5 = jnp.concatenate( [ input_embed_2, output_embed_2, input_embed_3, output_embed_3, input_embed_1, output_embed_1, ], axis=-1, ) context_6 = jnp.concatenate( [ input_embed_3, output_embed_3, input_embed_2, output_embed_2, input_embed_1, output_embed_1, ], axis=-1, ) context = jax.random.choice( key, jnp.array([context_1, context_2, context_3, context_4, context_5, context_6]), ) # Initialize state state = jnp.zeros((ds_size, channel_size)) state = state.at[..., 3 : 18 + 3].set(context) state = state.at[..., -10:].set(jax.nn.one_hot(sample[-1, 0], num_classes=10)) return state, sample[-1, -1] def sample_state(cs, key): """Sample state with data augmentation.""" key_sample, key_flip, key_perm, key_init = jax.random.split(key, 4) # Sample dataset _ = jax.random.choice(key_sample, train_tasks) sample = jax.random.choice(key_sample, train_ds) # Flip sample half of the time flip = jax.random.bernoulli(key_flip, p=0.5) sample = jnp.where(flip < 0.5, sample, jnp.flip(sample, axis=-1)) # Permute colors color_perm = jnp.concatenate( [ jnp.array([0], dtype=jnp.int32), jax.random.permutation(key_perm, jnp.arange(9)) + 1, ] ) sample = color_perm[sample] return create_state_with_sample(cs, sample, key_init) def sample_state_test(cs, key): """Sample state with data augmentation.""" key_sample, key_init = jax.random.split(key) # Sample dataset _ = jax.random.choice(key_sample, test_tasks) sample = jax.random.choice(key_sample, test_ds) return create_state_with_sample(cs, sample, key_init)

Instantiate system¶

In [10]:

class ARCNCA(ComplexSystem):
    """1D-ARC Neural Cellular Automata  class."""

    def __init__(self, *, rngs: nnx.Rngs):
        """Initialize 1D-ARC NCA."""
        self.perceive = ConvPerceive(
            channel_size=channel_size,
            perception_size=num_kernels * channel_size,
            kernel_size=(3,),
            feature_group_count=channel_size,
            rngs=rngs,
        )
        self.update = ResidualUpdate(
            num_spatial_dims=num_spatial_dims,
            channel_size=channel_size,
            perception_size=num_kernels * channel_size,
            hidden_layer_sizes=hidden_layer_sizes,
            cell_dropout_rate=cell_dropout_rate,
            zeros_init=True,
            rngs=rngs,
        )
        self.embed_input = nnx.Embed(
            num_embeddings=num_embeddings, features=features, rngs=rngs
        )

        # Initialize kernel with sobel filters
        kernel = jnp.concatenate(
            [
                identity_kernel(num_dims=num_spatial_dims),
                grad_kernel(num_dims=num_spatial_dims),
            ],
            axis=-1,
        )
        kernel = jnp.expand_dims(
            jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2
        )
        self.perceive.conv.kernel[...] = kernel

    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."""
        # Extract classification logits
        logits = state[..., -10:]

        # Render to RGB
        rgb = color_lookup[jnp.argmax(logits, axis=-1)]

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

class ARCNCA(ComplexSystem): """1D-ARC Neural Cellular Automata class.""" def __init__(self, *, rngs: nnx.Rngs): """Initialize 1D-ARC NCA.""" self.perceive = ConvPerceive( channel_size=channel_size, perception_size=num_kernels * channel_size, kernel_size=(3,), feature_group_count=channel_size, rngs=rngs, ) self.update = ResidualUpdate( num_spatial_dims=num_spatial_dims, channel_size=channel_size, perception_size=num_kernels * channel_size, hidden_layer_sizes=hidden_layer_sizes, cell_dropout_rate=cell_dropout_rate, zeros_init=True, rngs=rngs, ) self.embed_input = nnx.Embed( num_embeddings=num_embeddings, features=features, rngs=rngs ) # Initialize kernel with sobel filters kernel = jnp.concatenate( [ identity_kernel(num_dims=num_spatial_dims), grad_kernel(num_dims=num_spatial_dims), ], axis=-1, ) kernel = jnp.expand_dims( jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2 ) self.perceive.conv.kernel[...] = kernel 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.""" # Extract classification logits logits = state[..., -10:] # Render to RGB rgb = color_lookup[jnp.argmax(logits, axis=-1)] # Clip values to valid range and convert to uint8 return clip_and_uint8(rgb)

In [11]:

cs = ARCNCA(rngs=rngs)

cs = ARCNCA(rngs=rngs)

In [12]:

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: 49886

Train¶

In [13]:

num_train_steps = 100_000
lr_sched = optax.linear_schedule(
    init_value=learning_rate,
    end_value=0.1 * learning_rate,
    transition_steps=num_train_steps // 10,
)

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

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

num_train_steps = 100_000 lr_sched = optax.linear_schedule( init_value=learning_rate, end_value=0.1 * learning_rate, transition_steps=num_train_steps // 10, ) optimizer = optax.chain( optax.clip_by_global_norm(1.0), optax.adam(learning_rate=lr_sched), ) params = nnx.All(nnx.Param) optimizer = nnx.Optimizer(cs, optimizer, wrt=params)

Loss¶

In [14]:

def ce(state, output):
    """Cross-entropy."""
    return jnp.mean(
        optax.softmax_cross_entropy_with_integer_labels(state[..., -10:], output)
    )

def ce(state, output): """Cross-entropy.""" return jnp.mean( optax.softmax_cross_entropy_with_integer_labels(state[..., -10:], output) )

In [15]:

@nnx.jit
def loss_fn(cs, key):
    """Loss function."""
    keys = jax.random.split(key, batch_size)
    state, output = jax.vmap(sample_state, in_axes=(None, 0))(cs, keys)

    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 = ce(state, output)
    return loss

@nnx.jit def loss_fn(cs, key): """Loss function.""" keys = jax.random.split(key, batch_size) state, output = jax.vmap(sample_state, in_axes=(None, 0))(cs, keys) 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 = ce(state, output) return loss

Train step¶

In [16]:

@nnx.jit
def train_step(cs, optimizer, key):
    """Train step."""
    loss, grad = nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, params))(cs, key)
    optimizer.update(cs, grad)
    return loss

@nnx.jit def train_step(cs, optimizer, key): """Train step.""" loss, grad = nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, params))(cs, key) optimizer.update(cs, grad) return loss

Main loop¶

In [17]:

def accuracy(cs, eval_ds):
    """Compute accuracy."""
    eval_size = eval_ds.shape[0]
    state, output = jax.vmap(create_state_with_sample, in_axes=(None, 0, None))(
        cs, eval_ds, key
    )

    state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
    state = nnx.split_rngs(splits=eval_size)(
        nnx.vmap(
            lambda cs, state: cs(state, num_steps=num_steps),
            in_axes=(state_axes, 0),
        )
    )(cs, state)

    # Convert logits to symbols
    final_state_logits = state[..., -10:]
    final_state = jnp.argmax(final_state_logits, axis=-1)

    # Successful if all symbols match in the prediction
    return jnp.sum(jnp.all(final_state == output, axis=-1)) / eval_size

def accuracy(cs, eval_ds): """Compute accuracy.""" eval_size = eval_ds.shape[0] state, output = jax.vmap(create_state_with_sample, in_axes=(None, 0, None))( cs, eval_ds, key ) state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) state = nnx.split_rngs(splits=eval_size)( nnx.vmap( lambda cs, state: cs(state, num_steps=num_steps), in_axes=(state_axes, 0), ) )(cs, state) # Convert logits to symbols final_state_logits = state[..., -10:] final_state = jnp.argmax(final_state_logits, axis=-1) # Successful if all symbols match in the prediction return jnp.sum(jnp.all(final_state == output, axis=-1)) / eval_size

In [18]:

print_interval = 1_024

losses = []
start = time.perf_counter()
for i in range(num_train_steps):
    key, subkey = jax.random.split(key)
    loss = train_step(cs, optimizer, 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:])
        test_acc = accuracy(cs, test_ds)
        train_acc = accuracy(cs, train_ds)
        elapsed = time.perf_counter() - start
        print(
            f"Step {i:>6}/{num_train_steps} | {elapsed:6.1f}s | Loss {avg_loss:.3e} "
            f"| Test Acc {test_acc:.2%} | Train Acc {train_acc:.2%}"
        )

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

print_interval = 1_024 losses = [] start = time.perf_counter() for i in range(num_train_steps): key, subkey = jax.random.split(key) loss = train_step(cs, optimizer, 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:]) test_acc = accuracy(cs, test_ds) train_acc = accuracy(cs, train_ds) elapsed = time.perf_counter() - start print( f"Step {i:>6}/{num_train_steps} | {elapsed:6.1f}s | Loss {avg_loss:.3e} " f"| Test Acc {test_acc:.2%} | Train Acc {train_acc:.2%}" ) print(f"✨ Trained for {num_train_steps} steps in {time.perf_counter() - start:.0f}s")

Step      0/100000 |    6.0s | Loss 1.514e+00 | Test Acc 0.00% | Train Acc 0.00%

Step   1024/100000 |   14.8s | Loss 1.914e-01 | Test Acc 10.99% | Train Acc 10.37%

Step   2048/100000 |   23.7s | Loss 1.223e-01 | Test Acc 20.88% | Train Acc 21.60%

Step   3072/100000 |   32.5s | Loss 9.553e-02 | Test Acc 27.47% | Train Acc 33.70%

Step   4096/100000 |   41.4s | Loss 8.252e-02 | Test Acc 34.07% | Train Acc 41.23%

Step   5120/100000 |   50.3s | Loss 7.129e-02 | Test Acc 32.97% | Train Acc 42.59%

Step   6144/100000 |   59.3s | Loss 6.144e-02 | Test Acc 49.45% | Train Acc 55.43%

Step   7168/100000 |   68.1s | Loss 5.554e-02 | Test Acc 48.35% | Train Acc 60.49%

Step   8192/100000 |   77.1s | Loss 4.986e-02 | Test Acc 53.85% | Train Acc 61.11%

Step   9216/100000 |   86.0s | Loss 4.619e-02 | Test Acc 50.55% | Train Acc 63.95%

Step  10240/100000 |   94.9s | Loss 2.531e+00 | Test Acc 57.14% | Train Acc 65.80%

Step  11264/100000 |  104.0s | Loss 3.971e-02 | Test Acc 54.95% | Train Acc 67.16%

Step  12288/100000 |  113.1s | Loss 4.039e-02 | Test Acc 59.34% | Train Acc 67.78%

Step  13312/100000 |  122.0s | Loss 4.065e-02 | Test Acc 56.04% | Train Acc 68.02%

Step  14336/100000 |  131.0s | Loss 3.949e-02 | Test Acc 54.95% | Train Acc 68.27%

Step  15360/100000 |  140.1s | Loss 3.812e-02 | Test Acc 53.85% | Train Acc 69.14%

Step  16384/100000 |  149.1s | Loss 3.905e-02 | Test Acc 58.24% | Train Acc 69.26%

Step  17408/100000 |  158.1s | Loss 3.742e-02 | Test Acc 59.34% | Train Acc 69.51%

Step  18432/100000 |  167.2s | Loss 3.756e-02 | Test Acc 57.14% | Train Acc 69.88%

Step  19456/100000 |  176.2s | Loss 3.624e-02 | Test Acc 56.04% | Train Acc 71.11%

Step  20480/100000 |  185.3s | Loss 3.757e-02 | Test Acc 60.44% | Train Acc 72.10%

Step  21504/100000 |  194.3s | Loss 3.739e-02 | Test Acc 60.44% | Train Acc 70.49%

Step  22528/100000 |  203.3s | Loss 3.555e-02 | Test Acc 51.65% | Train Acc 71.48%

Step  23552/100000 |  212.3s | Loss 3.588e-02 | Test Acc 52.75% | Train Acc 71.60%

Step  24576/100000 |  221.4s | Loss 3.465e-02 | Test Acc 58.24% | Train Acc 73.09%

Step  25600/100000 |  230.4s | Loss 3.563e-02 | Test Acc 57.14% | Train Acc 72.47%

Step  26624/100000 |  239.4s | Loss 3.593e-02 | Test Acc 57.14% | Train Acc 74.94%

Step  27648/100000 |  248.4s | Loss 3.578e-02 | Test Acc 57.14% | Train Acc 73.58%

Step  28672/100000 |  257.4s | Loss 3.570e-02 | Test Acc 57.14% | Train Acc 70.37%

Step  29696/100000 |  266.4s | Loss 3.400e-02 | Test Acc 63.74% | Train Acc 71.60%

Step  30720/100000 |  275.4s | Loss 3.414e-02 | Test Acc 59.34% | Train Acc 74.20%

Step  31744/100000 |  284.4s | Loss 3.413e-02 | Test Acc 58.24% | Train Acc 72.35%

Step  32768/100000 |  293.4s | Loss 3.504e-02 | Test Acc 58.24% | Train Acc 74.07%

Step  33792/100000 |  302.6s | Loss 3.431e-02 | Test Acc 57.14% | Train Acc 73.33%

Step  34816/100000 |  311.6s | Loss 3.338e-02 | Test Acc 52.75% | Train Acc 74.07%

Step  35840/100000 |  320.7s | Loss 3.344e-02 | Test Acc 53.85% | Train Acc 72.96%

Step  36864/100000 |  329.7s | Loss 3.320e-02 | Test Acc 60.44% | Train Acc 73.21%

Step  37888/100000 |  338.8s | Loss 3.237e-02 | Test Acc 57.14% | Train Acc 73.09%

Step  38912/100000 |  347.8s | Loss 3.309e-02 | Test Acc 62.64% | Train Acc 75.06%

Step  39936/100000 |  356.9s | Loss 3.246e-02 | Test Acc 57.14% | Train Acc 74.07%

Step  40960/100000 |  365.9s | Loss 3.337e-02 | Test Acc 58.24% | Train Acc 74.69%

Step  41984/100000 |  374.9s | Loss 3.255e-02 | Test Acc 63.74% | Train Acc 75.06%

Step  43008/100000 |  384.2s | Loss 3.217e-02 | Test Acc 57.14% | Train Acc 75.80%

Step  44032/100000 |  393.3s | Loss 3.228e-02 | Test Acc 59.34% | Train Acc 75.31%

Step  45056/100000 |  402.3s | Loss 3.207e-02 | Test Acc 61.54% | Train Acc 74.32%

Step  46080/100000 |  411.3s | Loss 3.225e-02 | Test Acc 56.04% | Train Acc 75.06%

Step  47104/100000 |  420.5s | Loss 3.222e-02 | Test Acc 62.64% | Train Acc 76.67%

Step  48128/100000 |  429.5s | Loss 3.289e-02 | Test Acc 62.64% | Train Acc 75.31%

Step  49152/100000 |  438.6s | Loss 3.057e-02 | Test Acc 59.34% | Train Acc 73.83%

Step  50176/100000 |  447.7s | Loss 3.107e-02 | Test Acc 62.64% | Train Acc 75.19%

Step  51200/100000 |  456.9s | Loss 3.049e-02 | Test Acc 59.34% | Train Acc 76.67%

Step  52224/100000 |  465.9s | Loss 2.997e-02 | Test Acc 58.24% | Train Acc 75.68%

Step  53248/100000 |  475.0s | Loss 3.094e-02 | Test Acc 60.44% | Train Acc 76.30%

Step  54272/100000 |  484.1s | Loss 3.116e-02 | Test Acc 64.84% | Train Acc 75.80%

Step  55296/100000 |  493.2s | Loss 3.123e-02 | Test Acc 57.14% | Train Acc 74.69%

Step  56320/100000 |  502.4s | Loss 2.965e-02 | Test Acc 60.44% | Train Acc 76.05%

Step  57344/100000 |  511.5s | Loss 3.100e-02 | Test Acc 59.34% | Train Acc 76.05%

Step  58368/100000 |  520.7s | Loss 3.071e-02 | Test Acc 62.64% | Train Acc 75.43%

Step  59392/100000 |  529.7s | Loss 3.094e-02 | Test Acc 59.34% | Train Acc 73.83%

Step  60416/100000 |  538.9s | Loss 3.016e-02 | Test Acc 62.64% | Train Acc 75.31%

Step  61440/100000 |  548.1s | Loss 3.010e-02 | Test Acc 54.95% | Train Acc 73.83%

Step  62464/100000 |  557.3s | Loss 3.022e-02 | Test Acc 60.44% | Train Acc 76.17%

Step  63488/100000 |  566.4s | Loss 2.971e-02 | Test Acc 58.24% | Train Acc 75.06%

Step  64512/100000 |  575.6s | Loss 2.968e-02 | Test Acc 65.93% | Train Acc 76.91%

Step  65536/100000 |  584.7s | Loss 2.924e-02 | Test Acc 64.84% | Train Acc 77.16%

Step  66560/100000 |  593.9s | Loss 2.946e-02 | Test Acc 59.34% | Train Acc 75.93%

Step  67584/100000 |  603.1s | Loss 2.983e-02 | Test Acc 58.24% | Train Acc 74.94%

Step  68608/100000 |  612.2s | Loss 3.009e-02 | Test Acc 60.44% | Train Acc 75.31%

Step  69632/100000 |  621.5s | Loss 2.969e-02 | Test Acc 61.54% | Train Acc 77.28%

Step  70656/100000 |  630.7s | Loss 2.824e-02 | Test Acc 63.74% | Train Acc 76.54%

Step  71680/100000 |  639.8s | Loss 2.894e-02 | Test Acc 61.54% | Train Acc 72.35%

Step  72704/100000 |  648.9s | Loss 2.932e-02 | Test Acc 59.34% | Train Acc 75.56%

Step  73728/100000 |  658.2s | Loss 2.793e-02 | Test Acc 62.64% | Train Acc 76.05%

Step  74752/100000 |  667.3s | Loss 2.850e-02 | Test Acc 59.34% | Train Acc 75.68%

Step  75776/100000 |  676.5s | Loss 2.857e-02 | Test Acc 57.14% | Train Acc 75.80%

Step  76800/100000 |  685.6s | Loss 2.804e-02 | Test Acc 61.54% | Train Acc 77.04%

Step  77824/100000 |  694.8s | Loss 4.272e-02 | Test Acc 59.34% | Train Acc 76.42%

Step  78848/100000 |  704.2s | Loss 2.760e-02 | Test Acc 61.54% | Train Acc 77.41%

Step  79872/100000 |  713.4s | Loss 2.782e-02 | Test Acc 64.84% | Train Acc 77.16%

Step  80896/100000 |  722.6s | Loss 2.743e-02 | Test Acc 64.84% | Train Acc 76.17%

Step  81920/100000 |  732.0s | Loss 2.853e-02 | Test Acc 59.34% | Train Acc 76.79%

Step  82944/100000 |  741.3s | Loss 2.771e-02 | Test Acc 61.54% | Train Acc 75.93%

Step  83968/100000 |  750.5s | Loss 2.800e-02 | Test Acc 63.74% | Train Acc 76.91%

Step  84992/100000 |  759.7s | Loss 2.718e-02 | Test Acc 57.14% | Train Acc 76.05%

Step  86016/100000 |  768.8s | Loss 2.729e-02 | Test Acc 59.34% | Train Acc 76.30%

Step  87040/100000 |  778.0s | Loss 2.685e-02 | Test Acc 58.24% | Train Acc 75.68%

Step  88064/100000 |  787.2s | Loss 2.774e-02 | Test Acc 64.84% | Train Acc 76.05%

Step  89088/100000 |  796.5s | Loss 2.752e-02 | Test Acc 63.74% | Train Acc 77.65%

Step  90112/100000 |  805.7s | Loss 2.668e-02 | Test Acc 60.44% | Train Acc 76.54%

Step  91136/100000 |  814.8s | Loss 2.642e-02 | Test Acc 60.44% | Train Acc 77.41%

Step  92160/100000 |  824.3s | Loss 2.678e-02 | Test Acc 58.24% | Train Acc 76.79%

Step  93184/100000 |  833.5s | Loss 2.645e-02 | Test Acc 59.34% | Train Acc 76.67%

Step  94208/100000 |  842.6s | Loss 2.651e-02 | Test Acc 63.74% | Train Acc 77.04%

Step  95232/100000 |  851.9s | Loss 2.585e-02 | Test Acc 59.34% | Train Acc 75.80%

Step  96256/100000 |  861.3s | Loss 2.606e-02 | Test Acc 60.44% | Train Acc 75.68%

Step  97280/100000 |  870.4s | Loss 2.744e-02 | Test Acc 62.64% | Train Acc 78.02%

Step  98304/100000 |  879.6s | Loss 2.594e-02 | Test Acc 57.14% | Train Acc 78.40%

Step  99328/100000 |  888.8s | Loss 2.559e-02 | Test Acc 61.54% | Train Acc 77.16%

Step  99999/100000 |  894.8s | Loss 2.608e-02 | Test Acc 57.14% | Train Acc 75.68%
✨ Trained for 100000 steps in 895s

Run¶

In [19]:

num_examples = 8

key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, num_examples)
state_init, output = jax.vmap(sample_state_test, in_axes=(None, 0))(cs, 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=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, output = jax.vmap(sample_state_test, in_axes=(None, 0))(cs, 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=num_steps, return_states=True), in_axes=(state_axes, 0), ) )(cs, state_init)

Visualize¶

In [20]:

output_pred = jnp.argmax(state_final[..., -10:], axis=-1)
success = jnp.all(output_pred == output, axis=-1)

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

# Add titles to each image to indicate success
titles = ["Success" if s else "Failure" for s in success]
mediapy.show_images(frames, width=196, height=128, titles=titles)

output_pred = jnp.argmax(state_final[..., -10:], axis=-1) success = jnp.all(output_pred == output, axis=-1) states = jnp.concatenate([state_init[:, None], states], axis=1) frames = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, states) # Add titles to each image to indicate success titles = ["Success" if s else "Failure" for s in success] mediapy.show_images(frames, width=196, height=128, titles=titles)

Failure

Failure

Failure

Success

Failure

Failure

Failure

Success