15 类数据原型已能生成架构选择题

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Random frozen GRU language modeling: independent random token prefixes are continued by a single frozen, calibrated neural teacher. Predict next-token distributions only at supervised continuation positions. The teacher is shared across all splits; token targets of -100 are ignored.

### Synthesis (PyTorch)
```python
from __future__ import annotations
import hashlib
import math
from typing import Any
import torch
import torch.nn.functional as F

def _seed(seed: int, task: str, stream: str) -> int:
    return int.from_bytes(hashlib.sha256(f'{seed}|{task}|{stream}'.encode()).digest()[:8], 'little') % (2 ** 63 - 1)

class FrozenTeacher(torch.nn.Module):

    def __init__(self, config: dict, state: dict | None=None):
        super().__init__()
        self.net = TeacherNet(config['vocab_size'], config['teacher_width'], config['teacher_layers'])
        if state is not None:
            self.net.load_state_dict(state)
        self.scale = float(config['logit_scale']) / float(config['temperature'])
        self.requires_grad_(False)
        self.eval()

    def forward(self, x):
        return self.net(x) * self.scale

class TeacherNet(torch.nn.Module):

    def __init__(self, vocab, width, layers):
        super().__init__()
        self.token_embed = torch.nn.Embedding(vocab, width)
        self.gru = torch.nn.GRU(width, width, layers, batch_first=True)
        self.head = torch.nn.Linear(width, vocab)

    def forward(self, x):
        (h, _) = self.gru(self.token_embed(x))
        return self.head(h)

def _teacher_split(teacher, config, n, rng, prefix_length, continuation):
    seq = torch.randint(config['vocab_size'], (n, prefix_length), generator=rng)
    for _ in range(continuation):
        with torch.no_grad():
            probs = teacher(seq)[:, -1].softmax(-1)
        seq = torch.cat([seq, torch.multinomial(probs, 1, generator=rng)], 1)
    (x, y) = (seq[:, :-1], seq[:, 1:])
    with torch.no_grad():
        p = teacher(x).softmax(-1)
    mask = torch.zeros_like(x, dtype=torch.bool)
    mask[:, prefix_length - 1:] = True
    return {'x': x, 'y': y, 'mask': mask, 'teacher_probs': p}

def build(task_id: str, parameters: dict[str, Any], split_counts: dict[str, int], seed: int, data_seed: int | None=None):
    cfg = {'task_id': 'random_teacher_ntp'}
    (artifacts, metadata) = ({}, {})
    generator = torch.Generator().manual_seed(_seed(seed, 'random_teacher_ntp', 'generator'))
    teacher_cfg = parameters.get('teacher', {})
    prefix = int(parameters.get('prefix_length', min(parameters.get('prefix_lengths', [8]))))
    continuation = int(parameters.get('continuation_length', parameters.get('rollout_length', min(parameters.get('continuation_lengths', [16])))))
    cfg.update(vocab_size=int(parameters.get('vocab', 64)), prefix_length=prefix, continuation_length=continuation, rollout_length=continuation, context_length=prefix + 2 * continuation, teacher_width=int(teacher_cfg.get('width', 64)), teacher_layers=int(teacher_cfg.get('layers', 2)), temperature=float(teacher_cfg.get('temperature', 1.0)), logit_scale=float(teacher_cfg.get('logit_scale', parameters.get('logit_scale', 8.0))))
    with torch.random.fork_rng():
        torch.manual_seed(_seed(seed, 'random_teacher_ntp', 'teacher_weights'))
        teacher = FrozenTeacher(cfg)
    probe = torch.randint(cfg['vocab_size'], (128, int(parameters.get('teacher_probe_context', 8))), generator=generator)
    with torch.no_grad():
        for _ in range(12):
            p = teacher(probe).softmax(-1)[:, -1]
            entropy = float(-(p * p.clamp_min(1e-12).log()).sum(-1).mean())
            ratio = entropy / math.log(cfg['vocab_size'])
            if 0.35 <= ratio <= 0.88:
                break
            teacher.scale *= 1.5 if ratio > 0.88 else 0.7
        reset_p = teacher(probe[:, -1:]).softmax(-1)[:, -1]
        context_gap = float((p * (p.clamp_min(1e-12).log() - reset_p.clamp_min(1e-12).log())).sum(-1).mean())
    if not 0.35 <= ratio <= 0.88 or context_gap < 0.001:
        raise ValueError('Random teacher has insufficient entropy or context dependence')
    cfg['logit_scale'] = teacher.scale * cfg['temperature']
    artifacts['teacher_state'] = {k: v.detach().clone() for (k, v) in teacher.net.state_dict().items()}
    metadata.update(teacher_probe_entropy=entropy, teacher_context_kl_gap=context_gap, teacher_architecture='frozen_random_causal_gru', teacher_seed=_seed(seed, 'random_teacher_ntp', 'teacher_weights'), teacher_probe_scope='generator_only_not_test_selection')
    cfg.update(selection_metric='token_ce', higher_is_better=False)
    splits = {}
    seen_grammar = set()
    for (split, count) in split_counts.items():
        rng = torch.Generator().manual_seed(_seed(seed, 'random_teacher_ntp', split) if data_seed is None else data_seed)
        ood = split == 'ood'
        splits[split] = _teacher_split(teacher, cfg, count, rng, prefix, continuation * (2 if ood else 1))
    metadata.update(split_unit='whole_independently_sampled_sequence', seed_streams={s: _seed(seed, 'random_teacher_ntp', s) for s in ['generator', *split_counts]}, ood_axis={'random_teacher_ntp': 'continuation_length_x2', 'on_policy_distillation': 'continuation_length_x2', 'variable_order_markov': 'sequence_length_x2', 'copy_recall': 'distractor_length_x2', 'grammar_pretrain': 'terminal_frequency_shift', 'switching_state_space': 'switch_probability_x2', 'delay_dynamics': 'noise_scale_x1.8'}['random_teacher_ntp'])
    cfg['input_shape'] = list(splits['train']['x'].shape[1:])
    cfg['output_shape'] = list(splits['train']['y'].shape[1:])
    if 'vocab_size' in cfg:
        cfg['max_sequence_length'] = max((int(b['x'].shape[1]) for b in splits.values()))
    metadata['split_sizes'] = {s: len(b['x']) for (s, b) in splits.items()}
    return {'config': cfg, 'splits': splits, 'metadata': metadata, 'artifacts': artifacts}

def _canonical_parameters(params):
    p = dict(params)
    p['vocab'] = p.get('vocab_size', 64)
    p['sequence_length'] = p.get('context_length', 128) + (0 if p.get('task_id') == 'grammar_pretrain' else 1)
    p['teacher'] = {'width': p.get('teacher_width', 64), 'layers': p.get('teacher_layers', 2), 'temperature': p.get('temperature', 1.0), 'logit_scale': p.get('logit_scale', 8.0)}
    return p

def synthesize_task(name, params, seed, count):
    parameters = _canonical_parameters(dict(params, task_id='random_teacher_ntp'))
    dataset = build('random_teacher_ntp', parameters, {'train': count}, int(params['structure_seed']), data_seed=seed)
    batch = dataset['splits']['train']
    batch['y'] = batch['y'].clone()
    batch['y'][~batch['mask']] = -100
    return batch
TASK = 'random_teacher_ntp'
PARAMS = {'vocab_size': 64, 'prefix_length': 8, 'continuation_length': 16, 'context_length': 40, 'teacher_width': 64, 'teacher_layers': 2, 'temperature': 1.0, 'logit_scale': 8.0, 'train_size': 512, 'test_size': 128, 'validation_size': 128, 'ood_size': 128, 'instance_seed': 11, 'structure_seed': 10011, 'teacher_seed': 10011, 'ood_shift': {'prefix_length': 16, 'continuation_length': 8}, 'input_shape': [23], 'output_shape': [23], 'model_modes': ['bigram', 'ngram', 'gru', 'transformer'], 'objective': 'token_ce', 'output_kind': 'token'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_task(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 512 / 128 / 128
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: token_ce; model input/output shapes: [23] → [23].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 300
- batch_size: 16
- total_samples_seen: 4800 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_ce** on the held-out test split.
- Each choice is trained independently for **3** seeds (`0`..`2`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_ce across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: transformer
- width: 16
- depth: 1
- residual: True
- max_parameters: 2000000
- input_shape: [23]
- output_shape: [23]
- output_kind: token
- vocab_size: 64
- noise_steps: 100
- context_length: 40

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'transformer', 'width': 16, 'depth': 1, 'residual': True, 'max_parameters': 2000000, 'input_shape': [23], 'output_shape': [23], 'output_kind': 'token', 'vocab_size': 64, 'noise_steps': 100, 'context_length': 40}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class CausalTransformerLM(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        num_layers: int,
        num_heads: int,
        d_ff: int,
    ) -> None:
        super().__init__()
        self.context_length = context_length
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(context_length, d_model)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=num_heads,
            dim_feedforward=d_ff,
            dropout=0.0,
            batch_first=True,
            activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, seq_len = x.shape
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(b, -1)
        h = self.token_embed(x) + self.pos_embed(positions)
        mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
        h = self.encoder(h, mask=mask)
        return self.head(h)


class TokenModel(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        return self.model(x)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: transformer
- width: 48
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [23]
- output_shape: [23]
- output_kind: token
- vocab_size: 64
- noise_steps: 100
- context_length: 40

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'transformer', 'width': 48, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [23], 'output_shape': [23], 'output_kind': 'token', 'vocab_size': 64, 'noise_steps': 100, 'context_length': 40}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class CausalTransformerLM(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        num_layers: int,
        num_heads: int,
        d_ff: int,
    ) -> None:
        super().__init__()
        self.context_length = context_length
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(context_length, d_model)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=num_heads,
            dim_feedforward=d_ff,
            dropout=0.0,
            batch_first=True,
            activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, seq_len = x.shape
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(b, -1)
        h = self.token_embed(x) + self.pos_embed(positions)
        mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
        h = self.encoder(h, mask=mask)
        return self.head(h)


class TokenModel(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        return self.model(x)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: ngram
- width: 32
- depth: 3
- residual: False
- max_parameters: 2000000
- input_shape: [23]
- output_shape: [23]
- output_kind: token
- vocab_size: 64
- noise_steps: 100
- context_length: 40

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'ngram', 'width': 32, 'depth': 3, 'residual': False, 'max_parameters': 2000000, 'input_shape': [23], 'output_shape': [23], 'output_kind': 'token', 'vocab_size': 64, 'noise_steps': 100, 'context_length': 40}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class NgramLM(nn.Module):
    def __init__(self, vocab: int, width: int, order: int, depth: int = 1,
                 residual: bool = False) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocab, width)
        self.mix = nn.ModuleList(nn.Conv1d(width, width, order) for _ in range(depth))
        self.head = nn.Linear(width, vocab)
        self.order = order
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = self.embedding(x).transpose(1, 2)
        for layer in self.mix:
            z = F.gelu(layer(F.pad(h, (self.order - 1, 0))))
            h = h + z if self.residual else z
        return self.head(h.transpose(1, 2))


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Variable-order Markov language modeling: a fixed mixture of transition tables produces token sequences. Predict the next token at supervised positions. The tables are shared across splits; -100 targets are ignored.

### Synthesis (PyTorch)
```python
from __future__ import annotations
import hashlib
import math
from typing import Any
import torch
import torch.nn.functional as F

def _seed(seed: int, task: str, stream: str) -> int:
    return int.from_bytes(hashlib.sha256(f'{seed}|{task}|{stream}'.encode()).digest()[:8], 'little') % (2 ** 63 - 1)

def _markov_split(tables, config, n, rng, length):
    vocab = config['vocab_size']
    seq = torch.randint(vocab, (n, max(config['orders'])), generator=rng)
    all_probs = []
    for _ in range(length - seq.shape[1]):
        p = torch.zeros(n, vocab)
        for (order, weight, table) in zip(config['orders'], config['order_weights'], tables):
            idx = torch.zeros(n, dtype=torch.long)
            for token in seq[:, -order:].T:
                idx = idx * vocab + token
            p += weight * table[idx]
        all_probs.append(p)
        seq = torch.cat([seq, torch.multinomial(p, 1, generator=rng)], 1)
    (x, y) = (seq[:, :-1], seq[:, 1:])
    mask = torch.zeros_like(x, dtype=torch.bool)
    start = max(config['orders']) - 1
    mask[:, start:] = True
    p = torch.full((*x.shape, vocab), 1.0 / vocab)
    p[:, start:] = torch.stack(all_probs, 1)
    return {'x': x, 'y': y, 'mask': mask, 'teacher_probs': p}

def build(task_id: str, parameters: dict[str, Any], split_counts: dict[str, int], seed: int, data_seed: int | None=None):
    cfg = {'task_id': 'variable_order_markov'}
    (artifacts, metadata) = ({}, {})
    generator = torch.Generator().manual_seed(_seed(seed, 'variable_order_markov', 'generator'))
    orders = [int(o) for o in parameters.get('orders', [1, 2, 4])]
    weights = parameters.get('order_weights', [1 / len(orders)] * len(orders))
    if any((o <= 0 for o in orders)) or len(weights) != len(orders) or any((w < 0 for w in weights)) or (sum(weights) <= 0):
        raise ValueError('Invalid Markov orders or weights')
    cfg.update(vocab_size=int(parameters.get('vocab', 16)), orders=orders, order_weights=[w / sum(weights) for w in weights], context_length=int(parameters.get('sequence_length', 64)) - 1, selection_metric='token_ce', higher_is_better=False)
    tables = [(torch.randn(cfg['vocab_size'] ** o, cfg['vocab_size'], generator=generator) * 2).softmax(-1) for o in orders]
    artifacts['transition_tables'] = tables
    splits = {}
    seen_grammar = set()
    for (split, count) in split_counts.items():
        rng = torch.Generator().manual_seed(_seed(seed, 'variable_order_markov', split) if data_seed is None else data_seed)
        ood = split == 'ood'
        splits[split] = _markov_split(tables, cfg, count, rng, (cfg['context_length'] + 1) * (2 if ood else 1))
    metadata.update(split_unit='whole_independently_sampled_sequence', seed_streams={s: _seed(seed, 'variable_order_markov', s) for s in ['generator', *split_counts]}, ood_axis={'random_teacher_ntp': 'continuation_length_x2', 'on_policy_distillation': 'continuation_length_x2', 'variable_order_markov': 'sequence_length_x2', 'copy_recall': 'distractor_length_x2', 'grammar_pretrain': 'terminal_frequency_shift', 'switching_state_space': 'switch_probability_x2', 'delay_dynamics': 'noise_scale_x1.8'}['variable_order_markov'])
    cfg['input_shape'] = list(splits['train']['x'].shape[1:])
    cfg['output_shape'] = list(splits['train']['y'].shape[1:])
    if 'vocab_size' in cfg:
        cfg['max_sequence_length'] = max((int(b['x'].shape[1]) for b in splits.values()))
    metadata['split_sizes'] = {s: len(b['x']) for (s, b) in splits.items()}
    return {'config': cfg, 'splits': splits, 'metadata': metadata, 'artifacts': artifacts}

def _canonical_parameters(params):
    p = dict(params)
    p['vocab'] = p.get('vocab_size', 64)
    p['sequence_length'] = p.get('context_length', 128) + (0 if p.get('task_id') == 'grammar_pretrain' else 1)
    p['teacher'] = {'width': p.get('teacher_width', 64), 'layers': p.get('teacher_layers', 2), 'temperature': p.get('temperature', 1.0), 'logit_scale': p.get('logit_scale', 8.0)}
    return p

def synthesize_task(name, params, seed, count):
    parameters = _canonical_parameters(dict(params, task_id='variable_order_markov'))
    dataset = build('variable_order_markov', parameters, {'train': count}, int(params['structure_seed']), data_seed=seed)
    batch = dataset['splits']['train']
    batch['y'] = batch['y'].clone()
    batch['y'][~batch['mask']] = -100
    return batch
TASK = 'variable_order_markov'
PARAMS = {'vocab_size': 16, 'orders': [1, 2, 4], 'context_length': 63, 'train_size': 512, 'test_size': 128, 'validation_size': 128, 'ood_size': 128, 'instance_seed': 11, 'structure_seed': 10011, 'teacher_seed': 10011, 'ood_shift': {'orders': [1, 3], 'order_weights': [0.5, 0.5]}, 'input_shape': [63], 'output_shape': [63], 'model_modes': ['bigram', 'ngram', 'gru', 'transformer'], 'objective': 'token_ce', 'output_kind': 'token'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_task(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 512 / 128 / 128
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: token_ce; model input/output shapes: [63] → [63].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 300
- batch_size: 16
- total_samples_seen: 4800 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_ce** on the held-out test split.
- Each choice is trained independently for **3** seeds (`0`..`2`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_ce across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: transformer
- width: 32
- depth: 2
- residual: True
- max_parameters: 2000000
- input_shape: [63]
- output_shape: [63]
- output_kind: token
- vocab_size: 16
- noise_steps: 100
- context_length: 63

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'transformer', 'width': 32, 'depth': 2, 'residual': True, 'max_parameters': 2000000, 'input_shape': [63], 'output_shape': [63], 'output_kind': 'token', 'vocab_size': 16, 'noise_steps': 100, 'context_length': 63}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class CausalTransformerLM(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        num_layers: int,
        num_heads: int,
        d_ff: int,
    ) -> None:
        super().__init__()
        self.context_length = context_length
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(context_length, d_model)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=num_heads,
            dim_feedforward=d_ff,
            dropout=0.0,
            batch_first=True,
            activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, seq_len = x.shape
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(b, -1)
        h = self.token_embed(x) + self.pos_embed(positions)
        mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
        h = self.encoder(h, mask=mask)
        return self.head(h)


class TokenModel(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        return self.model(x)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: bigram
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [63]
- output_shape: [63]
- output_kind: token
- vocab_size: 16
- noise_steps: 100
- context_length: 63

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'bigram', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [63], 'output_shape': [63], 'output_kind': 'token', 'vocab_size': 16, 'noise_steps': 100, 'context_length': 63}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class NgramLM(nn.Module):
    def __init__(self, vocab: int, width: int, order: int, depth: int = 1,
                 residual: bool = False) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocab, width)
        self.mix = nn.ModuleList(nn.Conv1d(width, width, order) for _ in range(depth))
        self.head = nn.Linear(width, vocab)
        self.order = order
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = self.embedding(x).transpose(1, 2)
        for layer in self.mix:
            z = F.gelu(layer(F.pad(h, (self.order - 1, 0))))
            h = h + z if self.residual else z
        return self.head(h.transpose(1, 2))


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: ngram
- width: 32
- depth: 3
- residual: False
- max_parameters: 2000000
- input_shape: [63]
- output_shape: [63]
- output_kind: token
- vocab_size: 16
- noise_steps: 100
- context_length: 63

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'ngram', 'width': 32, 'depth': 3, 'residual': False, 'max_parameters': 2000000, 'input_shape': [63], 'output_shape': [63], 'output_kind': 'token', 'vocab_size': 16, 'noise_steps': 100, 'context_length': 63}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class NgramLM(nn.Module):
    def __init__(self, vocab: int, width: int, order: int, depth: int = 1,
                 residual: bool = False) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocab, width)
        self.mix = nn.ModuleList(nn.Conv1d(width, width, order) for _ in range(depth))
        self.head = nn.Linear(width, vocab)
        self.order = order
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = self.embedding(x).transpose(1, 2)
        for layer in self.mix:
            z = F.gelu(layer(F.pad(h, (self.order - 1, 0))))
            h = h + z if self.residual else z
        return self.head(h.transpose(1, 2))


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Delayed copy recall: an input sequence contains a source string, distracting tokens and a query marker. Learn to reproduce the entire source string. Training uses masked next-token cross-entropy. Ranking uses whole-string error under free autoregressive rollout, not teacher-forced token accuracy.

### Synthesis (PyTorch)
```python
from __future__ import annotations
import hashlib
import math
from typing import Any
import torch
import torch.nn.functional as F

def _seed(seed: int, task: str, stream: str) -> int:
    return int.from_bytes(hashlib.sha256(f'{seed}|{task}|{stream}'.encode()).digest()[:8], 'little') % (2 ** 63 - 1)

def _copy_split(config, n, rng, distractor_length):
    (v, k) = (config['vocab_size'], config['source_length'])
    source = torch.randint(3, v, (n, k), generator=rng)
    distract = torch.randint(3, v, (n, distractor_length), generator=rng)
    seq = torch.cat([torch.full((n, 1), 1), source, torch.full((n, 1), 2), distract, torch.full((n, 1), 1), source], 1)
    (x, y) = (seq[:, :-1], seq[:, 1:])
    mask = torch.zeros_like(x, dtype=torch.bool)
    mask[:, -k:] = True
    return {'x': x, 'y': y, 'mask': mask}

def build(task_id: str, parameters: dict[str, Any], split_counts: dict[str, int], seed: int, data_seed: int | None=None):
    cfg = {'task_id': 'copy_recall'}
    (artifacts, metadata) = ({}, {})
    generator = torch.Generator().manual_seed(_seed(seed, 'copy_recall', 'generator'))
    source = int(parameters.get('source_length', 16))
    distractor = int(parameters.get('distractor_length', 64))
    cfg.update(vocab_size=int(parameters.get('vocab', 32)), source_length=source, distractor_length=distractor, answer_start=source + distractor + 2, context_length=2 * source + 2 * distractor + 2, selection_metric='exact_match', higher_is_better=True)
    splits = {}
    seen_grammar = set()
    for (split, count) in split_counts.items():
        rng = torch.Generator().manual_seed(_seed(seed, 'copy_recall', split) if data_seed is None else data_seed)
        ood = split == 'ood'
        splits[split] = _copy_split(cfg, count, rng, cfg['distractor_length'] * (2 if ood else 1))
    metadata.update(split_unit='whole_independently_sampled_sequence', seed_streams={s: _seed(seed, 'copy_recall', s) for s in ['generator', *split_counts]}, ood_axis={'random_teacher_ntp': 'continuation_length_x2', 'on_policy_distillation': 'continuation_length_x2', 'variable_order_markov': 'sequence_length_x2', 'copy_recall': 'distractor_length_x2', 'grammar_pretrain': 'terminal_frequency_shift', 'switching_state_space': 'switch_probability_x2', 'delay_dynamics': 'noise_scale_x1.8'}['copy_recall'])
    cfg['input_shape'] = list(splits['train']['x'].shape[1:])
    cfg['output_shape'] = list(splits['train']['y'].shape[1:])
    if 'vocab_size' in cfg:
        cfg['max_sequence_length'] = max((int(b['x'].shape[1]) for b in splits.values()))
    metadata['split_sizes'] = {s: len(b['x']) for (s, b) in splits.items()}
    return {'config': cfg, 'splits': splits, 'metadata': metadata, 'artifacts': artifacts}

def _canonical_parameters(params):
    p = dict(params)
    p['vocab'] = p.get('vocab_size', 64)
    p['sequence_length'] = p.get('context_length', 128) + (0 if p.get('task_id') == 'grammar_pretrain' else 1)
    p['teacher'] = {'width': p.get('teacher_width', 64), 'layers': p.get('teacher_layers', 2), 'temperature': p.get('temperature', 1.0), 'logit_scale': p.get('logit_scale', 8.0)}
    return p

def synthesize_task(name, params, seed, count):
    parameters = _canonical_parameters(dict(params, task_id='copy_recall'))
    dataset = build('copy_recall', parameters, {'train': count}, int(params['structure_seed']), data_seed=seed)
    batch = dataset['splits']['train']
    batch['y'] = batch['y'].clone()
    batch['y'][~batch['mask']] = -100
    return batch
TASK = 'copy_recall'
PARAMS = {'vocab_size': 8, 'source_length': 4, 'distractor_length': 8, 'context_length': 130, 'train_size': 2048, 'test_size': 256, 'validation_size': 128, 'ood_size': 64, 'instance_seed': 311, 'structure_seed': 10311, 'teacher_seed': 10311, 'ood_shift': {'source_length': 8, 'distractor_length': 80}, 'input_shape': [18], 'output_shape': [18], 'model_modes': ['ngram', 'gru', 'transformer'], 'objective': 'token_ce', 'output_kind': 'token'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_task(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 2048 / 128 / 256
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: token_ce; model input/output shapes: [18] → [18].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 1000
- batch_size: 16
- total_samples_seen: 16000 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_error** on the held-out test split.
- Each choice is trained independently for **3** seeds (`10`..`12`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_error across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: ngram
- width: 64
- depth: 2
- residual: True
- max_parameters: 2000000
- input_shape: [18]
- output_shape: [18]
- output_kind: token
- vocab_size: 8
- context_length: 130

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'ngram', 'width': 64, 'depth': 2, 'residual': True, 'max_parameters': 2000000, 'input_shape': [18], 'output_shape': [18], 'output_kind': 'token', 'vocab_size': 8, 'context_length': 130}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class NgramLM(nn.Module):
    def __init__(self, vocab: int, width: int, order: int, depth: int = 1,
                 residual: bool = False) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocab, width)
        self.mix = nn.ModuleList(nn.Conv1d(width, width, order) for _ in range(depth))
        self.head = nn.Linear(width, vocab)
        self.order = order
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = self.embedding(x).transpose(1, 2)
        for layer in self.mix:
            z = F.gelu(layer(F.pad(h, (self.order - 1, 0))))
            h = h + z if self.residual else z
        return self.head(h.transpose(1, 2))


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: bigram
- width: 64
- depth: 2
- residual: True
- max_parameters: 2000000
- input_shape: [18]
- output_shape: [18]
- output_kind: token
- vocab_size: 8
- context_length: 130

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'bigram', 'width': 64, 'depth': 2, 'residual': True, 'max_parameters': 2000000, 'input_shape': [18], 'output_shape': [18], 'output_kind': 'token', 'vocab_size': 8, 'context_length': 130}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class NgramLM(nn.Module):
    def __init__(self, vocab: int, width: int, order: int, depth: int = 1,
                 residual: bool = False) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocab, width)
        self.mix = nn.ModuleList(nn.Conv1d(width, width, order) for _ in range(depth))
        self.head = nn.Linear(width, vocab)
        self.order = order
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = self.embedding(x).transpose(1, 2)
        for layer in self.mix:
            z = F.gelu(layer(F.pad(h, (self.order - 1, 0))))
            h = h + z if self.residual else z
        return self.head(h.transpose(1, 2))


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: gru
- width: 64
- depth: 2
- residual: True
- max_parameters: 2000000
- input_shape: [18]
- output_shape: [18]
- output_kind: token
- vocab_size: 8
- context_length: 130

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'gru', 'width': 64, 'depth': 2, 'residual': True, 'max_parameters': 2000000, 'input_shape': [18], 'output_shape': [18], 'output_kind': 'token', 'vocab_size': 8, 'context_length': 130}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class CausalGruLM(nn.Module):
    """A causal, unidirectional GRU language model."""

    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        num_layers: int,
        layer_residual: bool = False,
    ) -> None:
        super().__init__()
        self.context_length = context_length
        self.layer_residual = bool(layer_residual)
        self.token_embed = nn.Embedding(vocab_size, d_model)
        if self.layer_residual:
            # One independent single-layer GRU per architectural layer makes
            # the residual path explicit without adding any parameters.
            self.gru_layers = nn.ModuleList(
                [
                    nn.GRU(
                        input_size=d_model,
                        hidden_size=d_model,
                        num_layers=1,
                        batch_first=True,
                        dropout=0.0,
                    )
                    for _ in range(num_layers)
                ]
            )
        else:
            self.gru = nn.GRU(
                input_size=d_model,
                hidden_size=d_model,
                num_layers=num_layers,
                batch_first=True,
                dropout=0.0,
            )
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.token_embed(x)
        if self.layer_residual:
            for layer in self.gru_layers:
                layer_output, _ = layer(h)
                h = h + layer_output
        else:
            h, _ = self.gru(h)
        return self.head(h)


class TokenModel(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        return self.model(x)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Small language-model pretraining from scratch on strings generated by a frozen hierarchical grammar. Predict next tokens. Padding and unsupervised positions have target -100 and are ignored. Complete sequences are assigned to splits without exact sequence overlap.

### Synthesis (PyTorch)
```python
from __future__ import annotations
import hashlib
import math
from typing import Any
import torch
import torch.nn.functional as F

def _seed(seed: int, task: str, stream: str) -> int:
    return int.from_bytes(hashlib.sha256(f'{seed}|{task}|{stream}'.encode()).digest()[:8], 'little') % (2 ** 63 - 1)

def _grammar_split(config, n, rng, depth, seen, terminal_shift=False):

    def expr(d, force_branch=False):
        if d == 0 or (not force_branch and torch.rand((), generator=rng) < 0.28):
            return [int(torch.randint(7, config['vocab_size'], (), generator=rng))]
        op = int(torch.randint(5, 7, (), generator=rng))
        return [3] + expr(d - 1) + [op] + expr(d - 1) + [4]
    rows = []
    attempts = 0
    while len(rows) < n:
        attempts += 1
        if attempts > max(10000, n * 100):
            raise ValueError('Grammar support exhausted for requested unique split size')
        body = expr(depth, force_branch=True)
        if terminal_shift:
            body = [7 + (t - 7) % 5 if t >= 7 else t for t in body]
        if len(body) + 2 > config['context_length'] + 1:
            raise ValueError('Grammar context_length must accommodate maximum depth')
        key = tuple(body)
        if key not in seen:
            rows.append([2] + body + [1])
            seen.add(key)
    seq = torch.zeros(n, config['context_length'] + 1, dtype=torch.long)
    for (i, row) in enumerate(rows):
        seq[i, :len(row)] = torch.tensor(row)
    return {'x': seq[:, :-1], 'y': seq[:, 1:], 'mask': seq[:, 1:] != 0}

def build(task_id: str, parameters: dict[str, Any], split_counts: dict[str, int], seed: int, data_seed: int | None=None):
    cfg = {'task_id': 'grammar_pretrain'}
    (artifacts, metadata) = ({}, {})
    generator = torch.Generator().manual_seed(_seed(seed, 'grammar_pretrain', 'generator'))
    cfg.update(vocab_size=int(parameters.get('vocab', 32)), max_depth=int(parameters.get('max_depth', 4)), context_length=int(parameters.get('sequence_length', 128)), selection_metric='token_ce', higher_is_better=False)
    metadata['grammar'] = 'E -> terminal(.28) | (E+E) | (E*E); depth limit gives terminal'
    splits = {}
    seen_grammar = set()
    for (split, count) in split_counts.items():
        rng = torch.Generator().manual_seed(_seed(seed, 'grammar_pretrain', split) if data_seed is None else data_seed)
        ood = split == 'ood'
        splits[split] = _grammar_split(cfg, count, rng, cfg['max_depth'], seen_grammar, terminal_shift=ood)
    metadata.update(split_unit='whole_independently_sampled_sequence', seed_streams={s: _seed(seed, 'grammar_pretrain', s) for s in ['generator', *split_counts]}, ood_axis={'random_teacher_ntp': 'continuation_length_x2', 'on_policy_distillation': 'continuation_length_x2', 'variable_order_markov': 'sequence_length_x2', 'copy_recall': 'distractor_length_x2', 'grammar_pretrain': 'terminal_frequency_shift', 'switching_state_space': 'switch_probability_x2', 'delay_dynamics': 'noise_scale_x1.8'}['grammar_pretrain'])
    cfg['input_shape'] = list(splits['train']['x'].shape[1:])
    cfg['output_shape'] = list(splits['train']['y'].shape[1:])
    if 'vocab_size' in cfg:
        cfg['max_sequence_length'] = max((int(b['x'].shape[1]) for b in splits.values()))
    metadata['split_sizes'] = {s: len(b['x']) for (s, b) in splits.items()}
    return {'config': cfg, 'splits': splits, 'metadata': metadata, 'artifacts': artifacts}

def _canonical_parameters(params):
    p = dict(params)
    p['vocab'] = p.get('vocab_size', 64)
    p['sequence_length'] = p.get('context_length', 128) + (0 if p.get('task_id') == 'grammar_pretrain' else 1)
    p['teacher'] = {'width': p.get('teacher_width', 64), 'layers': p.get('teacher_layers', 2), 'temperature': p.get('temperature', 1.0), 'logit_scale': p.get('logit_scale', 8.0)}
    return p

def synthesize_task(name, params, seed, count):
    parameters = _canonical_parameters(dict(params, task_id='grammar_pretrain'))
    dataset = build('grammar_pretrain', parameters, {'train': count}, int(params['structure_seed']), data_seed=seed)
    batch = dataset['splits']['train']
    if 'grammar_pretrain' == 'grammar_pretrain' and 'instance_seed' in params:
        seen = set()
        for (index, split) in enumerate(('train', 'validation', 'test', 'ood')):
            prior_seed = params['instance_seed'] + 20000 + index * 10000
            if prior_seed >= seed:
                break
            _grammar_split(dataset['config'], params[split + '_size'], torch.Generator().manual_seed(prior_seed), int(params['max_depth']), seen)
        batch = _grammar_split(dataset['config'], count, torch.Generator().manual_seed(seed), int(params['max_depth']), seen)
    batch['y'] = batch['y'].clone()
    batch['y'][~batch['mask']] = -100
    return batch
TASK = 'grammar_pretrain'
PARAMS = {'vocab_size': 32, 'max_depth': 4, 'context_length': 128, 'train_size': 512, 'test_size': 128, 'validation_size': 128, 'ood_size': 128, 'instance_seed': 11, 'structure_seed': 10011, 'teacher_seed': 10011, 'ood_shift': {'max_depth': 3}, 'input_shape': [128], 'output_shape': [128], 'model_modes': ['bigram', 'ngram', 'gru', 'transformer'], 'objective': 'token_ce', 'output_kind': 'token'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_task(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 512 / 128 / 128
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: token_ce; model input/output shapes: [128] → [128].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 300
- batch_size: 16
- total_samples_seen: 4800 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_ce** on the held-out test split.
- Each choice is trained independently for **3** seeds (`0`..`2`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_ce across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: transformer
- width: 16
- depth: 2
- residual: True
- max_parameters: 2000000
- input_shape: [128]
- output_shape: [128]
- output_kind: token
- vocab_size: 32
- noise_steps: 100
- context_length: 128

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'transformer', 'width': 16, 'depth': 2, 'residual': True, 'max_parameters': 2000000, 'input_shape': [128], 'output_shape': [128], 'output_kind': 'token', 'vocab_size': 32, 'noise_steps': 100, 'context_length': 128}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class CausalTransformerLM(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        num_layers: int,
        num_heads: int,
        d_ff: int,
    ) -> None:
        super().__init__()
        self.context_length = context_length
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(context_length, d_model)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=num_heads,
            dim_feedforward=d_ff,
            dropout=0.0,
            batch_first=True,
            activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, seq_len = x.shape
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(b, -1)
        h = self.token_embed(x) + self.pos_embed(positions)
        mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
        h = self.encoder(h, mask=mask)
        return self.head(h)


class TokenModel(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        return self.model(x)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: bigram
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [128]
- output_shape: [128]
- output_kind: token
- vocab_size: 32
- noise_steps: 100
- context_length: 128

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'bigram', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [128], 'output_shape': [128], 'output_kind': 'token', 'vocab_size': 32, 'noise_steps': 100, 'context_length': 128}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class NgramLM(nn.Module):
    def __init__(self, vocab: int, width: int, order: int, depth: int = 1,
                 residual: bool = False) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocab, width)
        self.mix = nn.ModuleList(nn.Conv1d(width, width, order) for _ in range(depth))
        self.head = nn.Linear(width, vocab)
        self.order = order
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = self.embedding(x).transpose(1, 2)
        for layer in self.mix:
            z = F.gelu(layer(F.pad(h, (self.order - 1, 0))))
            h = h + z if self.residual else z
        return self.head(h.transpose(1, 2))


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: transformer
- width: 16
- depth: 1
- residual: True
- max_parameters: 2000000
- input_shape: [128]
- output_shape: [128]
- output_kind: token
- vocab_size: 32
- noise_steps: 100
- context_length: 128

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'transformer', 'width': 16, 'depth': 1, 'residual': True, 'max_parameters': 2000000, 'input_shape': [128], 'output_shape': [128], 'output_kind': 'token', 'vocab_size': 32, 'noise_steps': 100, 'context_length': 128}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class CausalTransformerLM(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        num_layers: int,
        num_heads: int,
        d_ff: int,
    ) -> None:
        super().__init__()
        self.context_length = context_length
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(context_length, d_model)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=num_heads,
            dim_feedforward=d_ff,
            dropout=0.0,
            batch_first=True,
            activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, seq_len = x.shape
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(b, -1)
        h = self.token_embed(x) + self.pos_embed(positions)
        mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
        h = self.encoder(h, mask=mask)
        return self.head(h)


class TokenModel(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        return self.model(x)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
On-policy distillation from scratch: the student samples continuations from its own current policy. A frozen neural teacher is queried on those student-generated prefixes. Optimize the declared divergence. Held-out ranking measures forward KL on supervised continuation positions of fixed evaluation sequences.

### Synthesis (PyTorch)
```python
from __future__ import annotations
import hashlib
import math
from typing import Any
import torch
import torch.nn.functional as F

def _seed(seed: int, task: str, stream: str) -> int:
    return int.from_bytes(hashlib.sha256(f'{seed}|{task}|{stream}'.encode()).digest()[:8], 'little') % (2 ** 63 - 1)

class FrozenTeacher(torch.nn.Module):

    def __init__(self, config: dict, state: dict | None=None):
        super().__init__()
        self.net = TeacherNet(config['vocab_size'], config['teacher_width'], config['teacher_layers'])
        if state is not None:
            self.net.load_state_dict(state)
        self.scale = float(config['logit_scale']) / float(config['temperature'])
        self.requires_grad_(False)
        self.eval()

    def forward(self, x):
        return self.net(x) * self.scale

class TeacherNet(torch.nn.Module):

    def __init__(self, vocab, width, layers):
        super().__init__()
        self.token_embed = torch.nn.Embedding(vocab, width)
        self.gru = torch.nn.GRU(width, width, layers, batch_first=True)
        self.head = torch.nn.Linear(width, vocab)

    def forward(self, x):
        (h, _) = self.gru(self.token_embed(x))
        return self.head(h)

def _teacher_split(teacher, config, n, rng, prefix_length, continuation):
    seq = torch.randint(config['vocab_size'], (n, prefix_length), generator=rng)
    for _ in range(continuation):
        with torch.no_grad():
            probs = teacher(seq)[:, -1].softmax(-1)
        seq = torch.cat([seq, torch.multinomial(probs, 1, generator=rng)], 1)
    (x, y) = (seq[:, :-1], seq[:, 1:])
    with torch.no_grad():
        p = teacher(x).softmax(-1)
    mask = torch.zeros_like(x, dtype=torch.bool)
    mask[:, prefix_length - 1:] = True
    return {'x': x, 'y': y, 'mask': mask, 'teacher_probs': p}

def build(task_id: str, parameters: dict[str, Any], split_counts: dict[str, int], seed: int, data_seed: int | None=None):
    cfg = {'task_id': 'on_policy_distillation'}
    (artifacts, metadata) = ({}, {})
    generator = torch.Generator().manual_seed(_seed(seed, 'on_policy_distillation', 'generator'))
    teacher_cfg = parameters.get('teacher', {})
    prefix = int(parameters.get('prefix_length', min(parameters.get('prefix_lengths', [8]))))
    continuation = int(parameters.get('continuation_length', parameters.get('rollout_length', min(parameters.get('continuation_lengths', [16])))))
    cfg.update(vocab_size=int(parameters.get('vocab', 64)), prefix_length=prefix, continuation_length=continuation, rollout_length=continuation, context_length=prefix + 2 * continuation, teacher_width=int(teacher_cfg.get('width', 64)), teacher_layers=int(teacher_cfg.get('layers', 2)), temperature=float(teacher_cfg.get('temperature', 1.0)), logit_scale=float(teacher_cfg.get('logit_scale', parameters.get('logit_scale', 8.0))))
    with torch.random.fork_rng():
        torch.manual_seed(_seed(seed, 'on_policy_distillation', 'teacher_weights'))
        teacher = FrozenTeacher(cfg)
    probe = torch.randint(cfg['vocab_size'], (128, int(parameters.get('teacher_probe_context', 8))), generator=generator)
    with torch.no_grad():
        for _ in range(12):
            p = teacher(probe).softmax(-1)[:, -1]
            entropy = float(-(p * p.clamp_min(1e-12).log()).sum(-1).mean())
            ratio = entropy / math.log(cfg['vocab_size'])
            if 0.35 <= ratio <= 0.88:
                break
            teacher.scale *= 1.5 if ratio > 0.88 else 0.7
        reset_p = teacher(probe[:, -1:]).softmax(-1)[:, -1]
        context_gap = float((p * (p.clamp_min(1e-12).log() - reset_p.clamp_min(1e-12).log())).sum(-1).mean())
    if not 0.35 <= ratio <= 0.88 or context_gap < 0.001:
        raise ValueError('Random teacher has insufficient entropy or context dependence')
    cfg['logit_scale'] = teacher.scale * cfg['temperature']
    artifacts['teacher_state'] = {k: v.detach().clone() for (k, v) in teacher.net.state_dict().items()}
    metadata.update(teacher_probe_entropy=entropy, teacher_context_kl_gap=context_gap, teacher_architecture='frozen_random_causal_gru', teacher_seed=_seed(seed, 'on_policy_distillation', 'teacher_weights'), teacher_probe_scope='generator_only_not_test_selection')
    cfg.update(selection_metric='forward_kl', higher_is_better=False)
    splits = {}
    seen_grammar = set()
    for (split, count) in split_counts.items():
        rng = torch.Generator().manual_seed(_seed(seed, 'on_policy_distillation', split) if data_seed is None else data_seed)
        ood = split == 'ood'
        splits[split] = _teacher_split(teacher, cfg, count, rng, prefix, continuation * (2 if ood else 1))
    metadata.update(split_unit='whole_independently_sampled_sequence', seed_streams={s: _seed(seed, 'on_policy_distillation', s) for s in ['generator', *split_counts]}, ood_axis={'random_teacher_ntp': 'continuation_length_x2', 'on_policy_distillation': 'continuation_length_x2', 'variable_order_markov': 'sequence_length_x2', 'copy_recall': 'distractor_length_x2', 'grammar_pretrain': 'terminal_frequency_shift', 'switching_state_space': 'switch_probability_x2', 'delay_dynamics': 'noise_scale_x1.8'}['on_policy_distillation'])
    cfg['input_shape'] = list(splits['train']['x'].shape[1:])
    cfg['output_shape'] = list(splits['train']['y'].shape[1:])
    if 'vocab_size' in cfg:
        cfg['max_sequence_length'] = max((int(b['x'].shape[1]) for b in splits.values()))
    metadata['split_sizes'] = {s: len(b['x']) for (s, b) in splits.items()}
    return {'config': cfg, 'splits': splits, 'metadata': metadata, 'artifacts': artifacts}

def _canonical_parameters(params):
    p = dict(params)
    p['vocab'] = p.get('vocab_size', 64)
    p['sequence_length'] = p.get('context_length', 128) + (0 if p.get('task_id') == 'grammar_pretrain' else 1)
    p['teacher'] = {'width': p.get('teacher_width', 64), 'layers': p.get('teacher_layers', 2), 'temperature': p.get('temperature', 1.0), 'logit_scale': p.get('logit_scale', 8.0)}
    return p

def synthesize_task(name, params, seed, count):
    parameters = _canonical_parameters(dict(params, task_id='on_policy_distillation'))
    dataset = build('on_policy_distillation', parameters, {'train': count}, int(params['structure_seed']), data_seed=seed)
    batch = dataset['splits']['train']
    batch['y'] = batch['y'].clone()
    batch['y'][~batch['mask']] = -100
    return batch
TASK = 'on_policy_distillation'
PARAMS = {'vocab_size': 64, 'prefix_length': 8, 'continuation_length': 16, 'opd_rollout_length': 4, 'context_length': 40, 'teacher_width': 64, 'teacher_layers': 2, 'temperature': 1.0, 'logit_scale': 8.0, 'train_size': 512, 'test_size': 128, 'validation_size': 128, 'ood_size': 128, 'instance_seed': 11, 'structure_seed': 10011, 'teacher_seed': 10011, 'ood_shift': {'prefix_length': 16, 'continuation_length': 8}, 'input_shape': [23], 'output_shape': [23], 'model_modes': ['gru', 'transformer'], 'objective': 'opd', 'output_kind': 'token'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_task(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 512 / 128 / 128
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: opd; model input/output shapes: [23] → [23].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 300
- batch_size: 16
- total_samples_seen: 4800 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_kl** on the held-out test split.
- Each choice is trained independently for **3** seeds (`0`..`2`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_kl across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: transformer
- width: 16
- depth: 1
- residual: True
- max_parameters: 2000000
- input_shape: [23]
- output_shape: [23]
- output_kind: token
- vocab_size: 64
- noise_steps: 100
- context_length: 40

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'transformer', 'width': 16, 'depth': 1, 'residual': True, 'max_parameters': 2000000, 'input_shape': [23], 'output_shape': [23], 'output_kind': 'token', 'vocab_size': 64, 'noise_steps': 100, 'context_length': 40}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class CausalTransformerLM(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        num_layers: int,
        num_heads: int,
        d_ff: int,
    ) -> None:
        super().__init__()
        self.context_length = context_length
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(context_length, d_model)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=num_heads,
            dim_feedforward=d_ff,
            dropout=0.0,
            batch_first=True,
            activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, seq_len = x.shape
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(b, -1)
        h = self.token_embed(x) + self.pos_embed(positions)
        mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
        h = self.encoder(h, mask=mask)
        return self.head(h)


class TokenModel(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        return self.model(x)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: reverse_kl

```python
def loss_fn(model, pred, target):
    teacher = F.log_softmax(target, -1)
    student = F.log_softmax(pred, -1)
    mode = 'reverse_kl'
    if mode == "reverse_kl":
        return (student.exp() * (student - teacher)).sum(-1).mean()
    if mode == "js":
        mixture = ((teacher.exp() + student.exp()) / 2).clamp_min(1e-12).log()
        return ((teacher.exp() * (teacher - mixture)).sum(-1).mean() + (student.exp() * (student - mixture)).sum(-1).mean()) / 2
    return (teacher.exp() * (teacher - student)).sum(-1).mean()
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: gru
- width: 32
- depth: 3
- residual: False
- max_parameters: 2000000
- input_shape: [23]
- output_shape: [23]
- output_kind: token
- vocab_size: 64
- noise_steps: 100
- context_length: 40

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'gru', 'width': 32, 'depth': 3, 'residual': False, 'max_parameters': 2000000, 'input_shape': [23], 'output_shape': [23], 'output_kind': 'token', 'vocab_size': 64, 'noise_steps': 100, 'context_length': 40}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class CausalGruLM(nn.Module):
    """A causal, unidirectional GRU language model."""

    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        num_layers: int,
        layer_residual: bool = False,
    ) -> None:
        super().__init__()
        self.context_length = context_length
        self.layer_residual = bool(layer_residual)
        self.token_embed = nn.Embedding(vocab_size, d_model)
        if self.layer_residual:
            # One independent single-layer GRU per architectural layer makes
            # the residual path explicit without adding any parameters.
            self.gru_layers = nn.ModuleList(
                [
                    nn.GRU(
                        input_size=d_model,
                        hidden_size=d_model,
                        num_layers=1,
                        batch_first=True,
                        dropout=0.0,
                    )
                    for _ in range(num_layers)
                ]
            )
        else:
            self.gru = nn.GRU(
                input_size=d_model,
                hidden_size=d_model,
                num_layers=num_layers,
                batch_first=True,
                dropout=0.0,
            )
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.token_embed(x)
        if self.layer_residual:
            for layer in self.gru_layers:
                layer_output, _ = layer(h)
                h = h + layer_output
        else:
            h, _ = self.gru(h)
        return self.head(h)


class TokenModel(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        return self.model(x)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: reverse_kl

```python
def loss_fn(model, pred, target):
    teacher = F.log_softmax(target, -1)
    student = F.log_softmax(pred, -1)
    mode = 'reverse_kl'
    if mode == "reverse_kl":
        return (student.exp() * (student - teacher)).sum(-1).mean()
    if mode == "js":
        mixture = ((teacher.exp() + student.exp()) / 2).clamp_min(1e-12).log()
        return ((teacher.exp() * (teacher - mixture)).sum(-1).mean() + (student.exp() * (student - mixture)).sum(-1).mean()) / 2
    return (teacher.exp() * (teacher - student)).sum(-1).mean()
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: transformer
- width: 48
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [23]
- output_shape: [23]
- output_kind: token
- vocab_size: 64
- noise_steps: 100
- context_length: 40

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'transformer', 'width': 48, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [23], 'output_shape': [23], 'output_kind': 'token', 'vocab_size': 64, 'noise_steps': 100, 'context_length': 40}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class CausalTransformerLM(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        num_layers: int,
        num_heads: int,
        d_ff: int,
    ) -> None:
        super().__init__()
        self.context_length = context_length
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(context_length, d_model)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=num_heads,
            dim_feedforward=d_ff,
            dropout=0.0,
            batch_first=True,
            activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, seq_len = x.shape
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(b, -1)
        h = self.token_embed(x) + self.pos_embed(positions)
        mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
        h = self.encoder(h, mask=mask)
        return self.head(h)


class TokenModel(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        return self.model(x)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: reverse_kl

```python
def loss_fn(model, pred, target):
    teacher = F.log_softmax(target, -1)
    student = F.log_softmax(pred, -1)
    mode = 'reverse_kl'
    if mode == "reverse_kl":
        return (student.exp() * (student - teacher)).sum(-1).mean()
    if mode == "js":
        mixture = ((teacher.exp() + student.exp()) / 2).clamp_min(1e-12).log()
        return ((teacher.exp() * (teacher - mixture)).sum(-1).mean() + (student.exp() * (student - mixture)).sum(-1).mean()) / 2
    return (teacher.exp() * (teacher - student)).sum(-1).mean()
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Direct multi-horizon forecasting of a partially observed switching linear state-space process. Input is a context of continuous observations; target is the next horizon of observations. Hidden regimes switch during each trajectory. Rank by held-out MSE.

### Synthesis (PyTorch)
```python
import torch

def synthesize_continuous(name, params, seed, count):
    generator = torch.Generator().manual_seed(seed)
    structure = torch.Generator().manual_seed(int(params['structure_seed']))
    (dim, context, horizon) = (params['input_dim'], params['context_length'], params['horizon'])
    total = context + horizon
    regimes = params['regimes']
    matrices = torch.randn(regimes, dim, dim, generator=structure)
    matrices = matrices / torch.linalg.matrix_norm(matrices, ord=2)[:, None, None]
    identity_mix = float(params.get('transition_identity_mix', 0.0))
    if not 0 <= identity_mix <= 1:
        raise ValueError('transition_identity_mix must be in [0,1]')
    matrices = ((1 - identity_mix) * matrices + identity_mix * torch.eye(dim)) * float(params.get('transition_radius', 0.85))
    observation = torch.randn(dim, dim, generator=structure) / dim ** 0.5
    states = torch.randn(count, dim, generator=generator)
    regime = torch.randint(regimes, (count,), generator=generator)
    values = []
    for _ in range(total):
        change = torch.rand(count, generator=generator) < float(params.get('switch_probability', 0.08))
        regime = torch.where(change, torch.randint(regimes, (count,), generator=generator), regime)
        states = torch.bmm(matrices[regime], states[..., None]).squeeze(-1)
        states += float(params.get('process_noise', 0.1)) * torch.randn(count, dim, generator=generator)
        values.append(states @ observation.T + params['noise_std'] * torch.randn(count, dim, generator=generator))
    series = torch.stack(values, dim=1)
    return {'x': series[:, :context], 'y': series[:, context:]}
TASK = 'switching_state_space'
PARAMS = {'input_dim': 4, 'context_length': 16, 'horizon': 4, 'regimes': 3, 'noise_std': 0.005, 'train_size': 2048, 'test_size': 256, 'validation_size': 128, 'ood_size': 128, 'transition_identity_mix': 0.9, 'transition_radius': 0.95, 'switch_probability': 0.02, 'process_noise': 0.1, 'instance_seed': 131, 'structure_seed': 10131, 'teacher_seed': 10131, 'ood_shift': {'noise_std': 0.01}, 'input_shape': [16, 4], 'output_shape': [4, 4], 'model_modes': ['mlp', 'gru', 'tcn', 'transformer'], 'objective': 'regression', 'output_kind': 'regression'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_continuous(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 2048 / 128 / 256
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: regression; model input/output shapes: [16, 4] → [4, 4].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 1000
- batch_size: 16
- total_samples_seen: 16000 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_mse** on the held-out test split.
- Each choice is trained independently for **3** seeds (`0`..`2`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_mse across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: gru
- width: 48
- depth: 3
- residual: False
- max_parameters: 2000000
- input_shape: [16, 4]
- output_shape: [4, 4]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 16

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'gru', 'width': 48, 'depth': 3, 'residual': False, 'max_parameters': 2000000, 'input_shape': [16, 4], 'output_shape': [4, 4], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 16}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SequenceForecaster(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, kind: str, depth: int) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        dim = math.prod(input_shape[1:])
        self.in_proj = nn.Linear(dim, width)
        if kind == 'gru_forecaster':
            self.body = nn.GRU(width, width, depth, batch_first=True)
        elif kind == 'transformer_forecaster':
            self.pos = nn.Parameter(torch.randn(1, input_shape[0], width) * 0.02)
            block = nn.TransformerEncoderLayer(width, 4, width * 2, dropout=0,
                                               batch_first=True, activation='gelu')
            self.body = nn.TransformerEncoder(block, depth)
        else:
            self.body = nn.ModuleList(nn.Conv1d(width, width, 3, dilation=2 ** i)
                                      for i in range(depth))
        self.head = nn.Linear(width, math.prod(output_shape))

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = self.in_proj(x.float().flatten(2))
        if self.kind == 'gru_forecaster':
            h, _ = self.body(h)
        elif self.kind == 'transformer_forecaster':
            length = h.shape[1]
            mask = torch.ones(length, length, device=x.device).triu(1).bool()
            h = self.body(h + self.pos[:, :length], mask=mask)
        else:
            h = h.transpose(1, 2)
            for i, layer in enumerate(self.body):
                h = h + F.gelu(layer(F.pad(h, (2 * 2 ** i, 0))))
            h = h.transpose(1, 2)
        return self.head(h[:, -1]).reshape(x.shape[0], *self.output_shape)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: mlp
- width: 48
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [16, 4]
- output_shape: [4, 4]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 16

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'mlp', 'width': 48, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [16, 4], 'output_shape': [4, 4], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 16}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class FlatMLP(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, residual: bool = False) -> None:
        super().__init__()
        self.output_shape = output_shape
        self.in_proj = nn.Linear(math.prod(input_shape), width)
        self.layers = nn.ModuleList(nn.Linear(width, width) for _ in range(depth))
        self.out_proj = nn.Linear(width, math.prod(output_shape))
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = F.gelu(self.in_proj(x.float().flatten(1)))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        return self.out_proj(h).reshape(x.shape[0], *self.output_shape)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: tcn
- width: 48
- depth: 1
- residual: True
- max_parameters: 2000000
- input_shape: [16, 4]
- output_shape: [4, 4]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 16

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'tcn', 'width': 48, 'depth': 1, 'residual': True, 'max_parameters': 2000000, 'input_shape': [16, 4], 'output_shape': [4, 4], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 16}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SequenceForecaster(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, kind: str, depth: int) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        dim = math.prod(input_shape[1:])
        self.in_proj = nn.Linear(dim, width)
        if kind == 'gru_forecaster':
            self.body = nn.GRU(width, width, depth, batch_first=True)
        elif kind == 'transformer_forecaster':
            self.pos = nn.Parameter(torch.randn(1, input_shape[0], width) * 0.02)
            block = nn.TransformerEncoderLayer(width, 4, width * 2, dropout=0,
                                               batch_first=True, activation='gelu')
            self.body = nn.TransformerEncoder(block, depth)
        else:
            self.body = nn.ModuleList(nn.Conv1d(width, width, 3, dilation=2 ** i)
                                      for i in range(depth))
        self.head = nn.Linear(width, math.prod(output_shape))

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = self.in_proj(x.float().flatten(2))
        if self.kind == 'gru_forecaster':
            h, _ = self.body(h)
        elif self.kind == 'transformer_forecaster':
            length = h.shape[1]
            mask = torch.ones(length, length, device=x.device).triu(1).bool()
            h = self.body(h + self.pos[:, :length], mask=mask)
        else:
            h = h.transpose(1, 2)
            for i, layer in enumerate(self.body):
                h = h + F.gelu(layer(F.pad(h, (2 * 2 ** i, 0))))
            h = h.transpose(1, 2)
        return self.head(h[:, -1]).reshape(x.shape[0], *self.output_shape)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Direct multi-horizon forecasting of a delayed nonlinear dynamical system. Input is a continuous context covering the configured delay; target is the next horizon. The recurrence uses both the latest value and a distant past value. Rank by held-out MSE.

### Synthesis (PyTorch)
```python
import torch

def synthesize_continuous(name, params, seed, count):
    generator = torch.Generator().manual_seed(seed)
    structure = torch.Generator().manual_seed(int(params['structure_seed']))
    (dim, context, horizon) = (params['input_dim'], params['context_length'], params['horizon'])
    total = context + horizon
    delay = int(params['delay'])
    if context < delay:
        raise ValueError('context must cover the configured delay')
    initial = torch.randn(count, delay, dim, generator=generator)
    values = list(initial.unbind(dim=1))
    mixing = torch.randn(dim, dim, generator=structure)
    mixing = mixing / torch.linalg.matrix_norm(mixing, ord=2)
    for t in range(delay, total):
        values.append(0.4 * values[-1] + 0.55 * torch.tanh(values[t - delay] @ mixing.T))
    series = torch.stack(values, dim=1)
    return {'x': series[:, :context], 'y': series[:, context:]}
TASK = 'delay_dynamics'
PARAMS = {'input_dim': 4, 'context_length': 64, 'horizon': 16, 'delay': 32, 'train_size': 512, 'test_size': 128, 'validation_size': 128, 'ood_size': 128, 'instance_seed': 11, 'structure_seed': 10011, 'teacher_seed': 10011, 'ood_shift': {'delay': 16}, 'input_shape': [64, 4], 'output_shape': [16, 4], 'model_modes': ['mlp', 'gru', 'tcn', 'transformer'], 'objective': 'regression', 'output_kind': 'regression'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_continuous(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 512 / 128 / 128
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: regression; model input/output shapes: [64, 4] → [16, 4].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 300
- batch_size: 16
- total_samples_seen: 4800 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_mse** on the held-out test split.
- Each choice is trained independently for **3** seeds (`0`..`2`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_mse across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: mlp
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [64, 4]
- output_shape: [16, 4]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 64

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'mlp', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [64, 4], 'output_shape': [16, 4], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 64}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class FlatMLP(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, residual: bool = False) -> None:
        super().__init__()
        self.output_shape = output_shape
        self.in_proj = nn.Linear(math.prod(input_shape), width)
        self.layers = nn.ModuleList(nn.Linear(width, width) for _ in range(depth))
        self.out_proj = nn.Linear(width, math.prod(output_shape))
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = F.gelu(self.in_proj(x.float().flatten(1)))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        return self.out_proj(h).reshape(x.shape[0], *self.output_shape)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: transformer
- width: 16
- depth: 1
- residual: True
- max_parameters: 2000000
- input_shape: [64, 4]
- output_shape: [16, 4]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 64

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'transformer', 'width': 16, 'depth': 1, 'residual': True, 'max_parameters': 2000000, 'input_shape': [64, 4], 'output_shape': [16, 4], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 64}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SequenceForecaster(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, kind: str, depth: int) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        dim = math.prod(input_shape[1:])
        self.in_proj = nn.Linear(dim, width)
        if kind == 'gru_forecaster':
            self.body = nn.GRU(width, width, depth, batch_first=True)
        elif kind == 'transformer_forecaster':
            self.pos = nn.Parameter(torch.randn(1, input_shape[0], width) * 0.02)
            block = nn.TransformerEncoderLayer(width, 4, width * 2, dropout=0,
                                               batch_first=True, activation='gelu')
            self.body = nn.TransformerEncoder(block, depth)
        else:
            self.body = nn.ModuleList(nn.Conv1d(width, width, 3, dilation=2 ** i)
                                      for i in range(depth))
        self.head = nn.Linear(width, math.prod(output_shape))

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = self.in_proj(x.float().flatten(2))
        if self.kind == 'gru_forecaster':
            h, _ = self.body(h)
        elif self.kind == 'transformer_forecaster':
            length = h.shape[1]
            mask = torch.ones(length, length, device=x.device).triu(1).bool()
            h = self.body(h + self.pos[:, :length], mask=mask)
        else:
            h = h.transpose(1, 2)
            for i, layer in enumerate(self.body):
                h = h + F.gelu(layer(F.pad(h, (2 * 2 ** i, 0))))
            h = h.transpose(1, 2)
        return self.head(h[:, -1]).reshape(x.shape[0], *self.output_shape)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: transformer
- width: 32
- depth: 2
- residual: True
- max_parameters: 2000000
- input_shape: [64, 4]
- output_shape: [16, 4]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 64

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'transformer', 'width': 32, 'depth': 2, 'residual': True, 'max_parameters': 2000000, 'input_shape': [64, 4], 'output_shape': [16, 4], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 64}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SequenceForecaster(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, kind: str, depth: int) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        dim = math.prod(input_shape[1:])
        self.in_proj = nn.Linear(dim, width)
        if kind == 'gru_forecaster':
            self.body = nn.GRU(width, width, depth, batch_first=True)
        elif kind == 'transformer_forecaster':
            self.pos = nn.Parameter(torch.randn(1, input_shape[0], width) * 0.02)
            block = nn.TransformerEncoderLayer(width, 4, width * 2, dropout=0,
                                               batch_first=True, activation='gelu')
            self.body = nn.TransformerEncoder(block, depth)
        else:
            self.body = nn.ModuleList(nn.Conv1d(width, width, 3, dilation=2 ** i)
                                      for i in range(depth))
        self.head = nn.Linear(width, math.prod(output_shape))

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = self.in_proj(x.float().flatten(2))
        if self.kind == 'gru_forecaster':
            h, _ = self.body(h)
        elif self.kind == 'transformer_forecaster':
            length = h.shape[1]
            mask = torch.ones(length, length, device=x.device).triu(1).bool()
            h = self.body(h + self.pos[:, :length], mask=mask)
        else:
            h = h.transpose(1, 2)
            for i, layer in enumerate(self.body):
                h = h + F.gelu(layer(F.pad(h, (2 * 2 ** i, 0))))
            h = h.transpose(1, 2)
        return self.head(h[:, -1]).reshape(x.shape[0], *self.output_shape)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Synthetic image relation classification. Independently rendered geometric objects define relation labels. Models receive images and predict class logits. Rank by held-out classification error.

### Synthesis (PyTorch)
```python
from __future__ import annotations
import math
import hashlib
from typing import Any
import torch

def synthesize_task(name, params, seed, count):
    cfg = dict(params)
    cfg.setdefault('sensor_seed', int(cfg.get('structure_seed', 20261010)))
    batch = generate('shape_classification', count, seed, cfg)
    validate('shape_classification', batch, cfg)
    return batch

def render_shape(conditions: torch.Tensor, size: int=32) -> torch.Tensor:
    """Render [circle=0/square=1, center_x, center_y, radius, intensity]."""
    p = torch.arange(size, dtype=torch.float32)
    (yy, xx) = torch.meshgrid(p, p, indexing='ij')
    c = conditions[:, :, None, None]
    (dx, dy) = (xx - c[:, 1], yy - c[:, 2])
    circle = dx.square() + dy.square() <= c[:, 3].square()
    square = torch.maximum(dx.abs(), dy.abs()) <= c[:, 3]
    mask = torch.where(c[:, 0].bool(), square, circle)
    return (mask * c[:, 4]).float().unsqueeze(1)

def generate(task_id: str, n: int, seed: int, config: dict[str, Any] | None=None) -> dict[str, torch.Tensor]:
    """Generate whole independent examples. Caller owns split seed streams."""
    cfg = config or {}
    if n < 1:
        raise ValueError('n must be positive')
    rng = torch.Generator().manual_seed(seed)
    size = int(cfg.get('size', 32))
    if size < 16:
        raise ValueError('spatial grid must contain at least 16 points per side')
    rng = torch.Generator().manual_seed(seed)
    labels = torch.arange(n) % 4
    labels = labels[torch.randperm(n, generator=rng)]
    circle = torch.zeros(n, 5)
    square = torch.zeros(n, 5)
    square[:, 0] = 1
    circle[:, 4] = square[:, 4] = 1
    side_labels = labels < 2
    count = int(side_labels.sum())
    left = torch.rand(count, 2, generator=rng)
    center_left = torch.stack((size * (0.2 + 0.08 * left[:, 0]), size * (0.25 + 0.5 * left[:, 1])), 1)
    right = torch.rand(count, 2, generator=rng)
    center_right = torch.stack((size * (0.72 + 0.08 * right[:, 0]), size * (0.25 + 0.5 * right[:, 1])), 1)
    is_left = labels[side_labels] == 0
    circle[side_labels, 1:3] = torch.where(is_left[:, None], center_left, center_right)
    square[side_labels, 1:3] = torch.where(is_left[:, None], center_right, center_left)
    circle[side_labels, 3] = size * (0.055 + 0.025 * torch.rand(count, generator=rng))
    square[side_labels, 3] = size * (0.055 + 0.025 * torch.rand(count, generator=rng))
    nested = ~side_labels
    count = int(nested.sum())
    centers = size * (0.4 + 0.2 * torch.rand(count, 2, generator=rng))
    circle[nested, 1:3] = square[nested, 1:3] = centers
    outer = size * (0.17 + 0.03 * torch.rand(count, generator=rng))
    inner = size * (0.055 + 0.01 * torch.rand(count, generator=rng))
    circle_inside = labels[nested] == 2
    circle[nested, 3] = torch.where(circle_inside, inner, outer)
    square[nested, 3] = torch.where(circle_inside, outer, inner)
    offsets = (torch.rand(count, 2, generator=rng) * 2 - 1) * (size * 0.025)
    circle[nested, 1:3] += offsets * circle_inside[:, None]
    square[nested, 1:3] += offsets * (~circle_inside)[:, None]
    circle[:, 3] *= float(cfg.get('shape_radius_scale', 1.0))
    square[:, 3] *= float(cfg.get('shape_radius_scale', 1.0))
    image = torch.cat((render_shape(circle, size), render_shape(square, size)), 1)
    return {'x': image, 'y': labels.long(), 'geometry': torch.stack((circle, square), 1)}

def validate(task_id: str, batch: dict[str, torch.Tensor], config: dict[str, Any] | None=None) -> dict[str, float | bool]:
    cfg = config or {}
    (x, y) = (batch['x'], batch['y'])
    if not torch.isfinite(x).all() or not torch.isfinite(y).all():
        raise ValueError('Nonfinite spatial dataset')
    if x.shape[0] != y.shape[0] or x.shape[0] < 1:
        raise ValueError('Mismatched spatial example count')
    evidence: dict[str, float | bool] = {'finite': True, 'n': float(x.shape[0])}
    (c, s) = batch['geometry'].unbind(1)
    circle_inside = ((c[:, 1:3] - s[:, 1:3]).abs() + c[:, 3:4] <= s[:, 3:4]).all(1)
    delta = (c[:, 1:3] - s[:, 1:3]).abs() + s[:, 3:4]
    square_inside = delta.square().sum(1) <= c[:, 3].square()
    expected = torch.where(circle_inside, 2, torch.where(square_inside, 3, torch.where(c[:, 1] < s[:, 1], 0, 1)))
    if not torch.equal(expected, y):
        raise ValueError('Shape relations disagree with geometry')
    if not torch.equal(x, torch.cat((render_shape(c, x.shape[-1]), render_shape(s, x.shape[-1])), 1)):
        raise ValueError('Shape rendering disagrees with geometry')
    evidence['relation_verifier_accuracy'] = 1.0
    return evidence
TASK = 'shape_classification'
PARAMS = {'size': 32, 'num_classes': 4, 'train_size': 512, 'test_size': 128, 'validation_size': 128, 'ood_size': 128, 'instance_seed': 11, 'structure_seed': 10011, 'teacher_seed': 10011, 'ood_shift': {'shape_radius_scale': 1.15}, 'input_shape': [2, 32, 32], 'output_shape': [4], 'model_modes': ['mlp', 'cnn', 'residual_cnn'], 'objective': 'classification', 'output_kind': 'classification'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_task(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 512 / 128 / 128
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: classification; model input/output shapes: [2, 32, 32] → [4].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 300
- batch_size: 16
- total_samples_seen: 4800 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_error** on the held-out test split.
- Each choice is trained independently for **3** seeds (`0`..`2`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_error across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: cnn
- width: 32
- depth: 2
- residual: False
- max_parameters: 2000000
- input_shape: [2, 32, 32]
- output_shape: [4]
- output_kind: classification
- vocab_size: 1
- noise_steps: 100
- context_length: 2

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'cnn', 'width': 32, 'depth': 2, 'residual': False, 'max_parameters': 2000000, 'input_shape': [2, 32, 32], 'output_shape': [4], 'output_kind': 'classification', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 2}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SpatialModel(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, kind: str, condition_dim: int = 0,
                 residual: bool | None = None) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        self.condition_dim = condition_dim
        self.residual = kind in ('residual_cnn', 'spectral_conv') if residual is None else residual
        self.in_proj = nn.Conv2d(input_shape[0] + condition_dim, width, 3, padding=1)
        if kind == 'spectral_conv':
            self.layers = nn.ModuleList(SpectralLayer(width, 6) for _ in range(depth))
        else:
            self.layers = nn.ModuleList(nn.Conv2d(width, width, 3, padding=1) for _ in range(depth))
        classification = len(output_shape) == 1
        self.classification = classification
        self.head = nn.Linear(width, output_shape[0]) if classification else nn.Conv2d(width, output_shape[0], 1)

    def forward(self, x: torch.Tensor, condition: torch.Tensor | None = None,
                timestep: torch.Tensor | None = None, **_: Any) -> torch.Tensor:
        if self.condition_dim:
            if condition is None:
                raise ValueError('condition is required for this denoiser')
            cond = condition.float()
            if timestep is not None:
                cond = torch.cat((cond, timestep.reshape(-1, 1).float()), 1)
            if cond.shape[1] != self.condition_dim:
                raise ValueError('condition feature count does not match candidate spec')
            x = torch.cat((x, cond[:, :, None, None].expand(-1, -1, *x.shape[-2:])), 1)
        h = F.gelu(self.in_proj(x.float()))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        if self.classification:
            return self.head(h.mean((-2, -1)))
        return self.head(h)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: mlp
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [2, 32, 32]
- output_shape: [4]
- output_kind: classification
- vocab_size: 1
- noise_steps: 100
- context_length: 2

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'mlp', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [2, 32, 32], 'output_shape': [4], 'output_kind': 'classification', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 2}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class FlatMLP(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, residual: bool = False) -> None:
        super().__init__()
        self.output_shape = output_shape
        self.in_proj = nn.Linear(math.prod(input_shape), width)
        self.layers = nn.ModuleList(nn.Linear(width, width) for _ in range(depth))
        self.out_proj = nn.Linear(width, math.prod(output_shape))
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = F.gelu(self.in_proj(x.float().flatten(1)))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        return self.out_proj(h).reshape(x.shape[0], *self.output_shape)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: cnn
- width: 16
- depth: 1
- residual: True
- max_parameters: 2000000
- input_shape: [2, 32, 32]
- output_shape: [4]
- output_kind: classification
- vocab_size: 1
- noise_steps: 100
- context_length: 2

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'cnn', 'width': 16, 'depth': 1, 'residual': True, 'max_parameters': 2000000, 'input_shape': [2, 32, 32], 'output_shape': [4], 'output_kind': 'classification', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 2}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SpatialModel(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, kind: str, condition_dim: int = 0,
                 residual: bool | None = None) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        self.condition_dim = condition_dim
        self.residual = kind in ('residual_cnn', 'spectral_conv') if residual is None else residual
        self.in_proj = nn.Conv2d(input_shape[0] + condition_dim, width, 3, padding=1)
        if kind == 'spectral_conv':
            self.layers = nn.ModuleList(SpectralLayer(width, 6) for _ in range(depth))
        else:
            self.layers = nn.ModuleList(nn.Conv2d(width, width, 3, padding=1) for _ in range(depth))
        classification = len(output_shape) == 1
        self.classification = classification
        self.head = nn.Linear(width, output_shape[0]) if classification else nn.Conv2d(width, output_shape[0], 1)

    def forward(self, x: torch.Tensor, condition: torch.Tensor | None = None,
                timestep: torch.Tensor | None = None, **_: Any) -> torch.Tensor:
        if self.condition_dim:
            if condition is None:
                raise ValueError('condition is required for this denoiser')
            cond = condition.float()
            if timestep is not None:
                cond = torch.cat((cond, timestep.reshape(-1, 1).float()), 1)
            if cond.shape[1] != self.condition_dim:
                raise ValueError('condition feature count does not match candidate spec')
            x = torch.cat((x, cond[:, :, None, None].expand(-1, -1, *x.shape[-2:])), 1)
        h = F.gelu(self.in_proj(x.float()))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        if self.classification:
            return self.head(h.mean((-2, -1)))
        return self.head(h)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: cross-entropy on minibatch target labels

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if pred.ndim == 3:
        vocab = pred.shape[-1]
        return torch.nn.functional.cross_entropy(pred.reshape(-1, vocab), target.reshape(-1))
    return torch.nn.functional.cross_entropy(pred, target.reshape(-1))
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Conditional image generation with diffusion. Condition parameters specify the geometric object; training predicts sampled noise from noisy images, conditions and time. Evaluation starts from random noise and performs iterative generation. Rank by one minus independently verified condition satisfaction, including IoU at least 0.7, foreground intensity and background leakage checks.

### Synthesis (PyTorch)
```python
from __future__ import annotations
import math
import hashlib
from typing import Any
import torch

def synthesize_task(name, params, seed, count):
    cfg = dict(params)
    cfg.setdefault('sensor_seed', int(cfg.get('structure_seed', 20261010)))
    batch = generate('shape_diffusion', count, seed, cfg)
    validate('shape_diffusion', batch, cfg)
    return batch

def render_shape(conditions: torch.Tensor, size: int=32) -> torch.Tensor:
    """Render [circle=0/square=1, center_x, center_y, radius, intensity]."""
    p = torch.arange(size, dtype=torch.float32)
    (yy, xx) = torch.meshgrid(p, p, indexing='ij')
    c = conditions[:, :, None, None]
    (dx, dy) = (xx - c[:, 1], yy - c[:, 2])
    circle = dx.square() + dy.square() <= c[:, 3].square()
    square = torch.maximum(dx.abs(), dy.abs()) <= c[:, 3]
    mask = torch.where(c[:, 0].bool(), square, circle)
    return (mask * c[:, 4]).float().unsqueeze(1)

def generate(task_id: str, n: int, seed: int, config: dict[str, Any] | None=None) -> dict[str, torch.Tensor]:
    """Generate whole independent examples. Caller owns split seed streams."""
    cfg = config or {}
    if n < 1:
        raise ValueError('n must be positive')
    rng = torch.Generator().manual_seed(seed)
    size = int(cfg.get('size', 32))
    if size < 16:
        raise ValueError('spatial grid must contain at least 16 points per side')
    rng = torch.Generator().manual_seed(seed)
    conditions = torch.rand(n, 5, generator=rng)
    conditions[:, 0] = torch.randint(2, (n,), generator=rng).float()
    conditions[:, 1:3] = size * (0.3 + conditions[:, 1:3] * 0.4)
    if 'condition_position_margin' in cfg:
        margin = float(cfg['condition_position_margin'])
        if not 0.15 <= margin < 0.3:
            raise ValueError('OOD condition margin must keep all shapes inside image')
        conditions[:, 1:3] = torch.where(conditions[:, 1:3] < size / 2, size * margin, size * (1 - margin))
    conditions[:, 3] = size * (0.07 + conditions[:, 3] * 0.08)
    conditions[:, 4] = 0.65 + conditions[:, 4] * 0.35
    return {'x': conditions, 'y': render_shape(conditions, size)}

def validate(task_id: str, batch: dict[str, torch.Tensor], config: dict[str, Any] | None=None) -> dict[str, float | bool]:
    cfg = config or {}
    (x, y) = (batch['x'], batch['y'])
    if not torch.isfinite(x).all() or not torch.isfinite(y).all():
        raise ValueError('Nonfinite spatial dataset')
    if x.shape[0] != y.shape[0] or x.shape[0] < 1:
        raise ValueError('Mismatched spatial example count')
    evidence: dict[str, float | bool] = {'finite': True, 'n': float(x.shape[0])}
    if not torch.equal(y, render_shape(x, y.shape[-1])):
        raise ValueError('Conditional image target disagrees with renderer')
    evidence['renderer_error'] = 0.0
    return evidence
TASK = 'shape_diffusion'
PARAMS = {'size': 16, 'noise_steps': 32, 'train_size': 1024, 'test_size': 128, 'validation_size': 64, 'ood_size': 64, 'instance_seed': 101, 'structure_seed': 10101, 'teacher_seed': 10101, 'ood_shift': {'condition_position_margin': 0.2}, 'input_shape': [7, 16, 16], 'output_shape': [1, 16, 16], 'model_modes': ['cnn', 'residual_cnn', 'fno'], 'objective': 'diffusion', 'output_kind': 'diffusion', 'condition_dim': 5}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_task(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 1024 / 64 / 128
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: diffusion; model input/output shapes: [7, 16, 16] → [1, 16, 16].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 1000
- batch_size: 16
- total_samples_seen: 16000 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_condition_error** on the held-out test split.
- Each choice is trained independently for **3** seeds (`101`..`103`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_condition_error across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: residual_cnn
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [7, 16, 16]
- output_shape: [1, 16, 16]
- output_kind: diffusion
- vocab_size: 1
- noise_steps: 32
- context_length: 7

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'residual_cnn', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [7, 16, 16], 'output_shape': [1, 16, 16], 'output_kind': 'diffusion', 'vocab_size': 1, 'noise_steps': 32, 'context_length': 7}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SpatialModel(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, kind: str, condition_dim: int = 0,
                 residual: bool | None = None) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        self.condition_dim = condition_dim
        self.residual = kind in ('residual_cnn', 'spectral_conv') if residual is None else residual
        self.in_proj = nn.Conv2d(input_shape[0] + condition_dim, width, 3, padding=1)
        if kind == 'spectral_conv':
            self.layers = nn.ModuleList(SpectralLayer(width, 6) for _ in range(depth))
        else:
            self.layers = nn.ModuleList(nn.Conv2d(width, width, 3, padding=1) for _ in range(depth))
        classification = len(output_shape) == 1
        self.classification = classification
        self.head = nn.Linear(width, output_shape[0]) if classification else nn.Conv2d(width, output_shape[0], 1)

    def forward(self, x: torch.Tensor, condition: torch.Tensor | None = None,
                timestep: torch.Tensor | None = None, **_: Any) -> torch.Tensor:
        if self.condition_dim:
            if condition is None:
                raise ValueError('condition is required for this denoiser')
            cond = condition.float()
            if timestep is not None:
                cond = torch.cat((cond, timestep.reshape(-1, 1).float()), 1)
            if cond.shape[1] != self.condition_dim:
                raise ValueError('condition feature count does not match candidate spec')
            x = torch.cat((x, cond[:, :, None, None].expand(-1, -1, *x.shape[-2:])), 1)
        h = F.gelu(self.in_proj(x.float()))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        if self.classification:
            return self.head(h.mean((-2, -1)))
        return self.head(h)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.001
- Weight decay: 0.0

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: diffusion_mse

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: coordinate_cnn
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [7, 16, 16]
- output_shape: [1, 16, 16]
- output_kind: diffusion
- vocab_size: 1
- noise_steps: 32
- context_length: 7

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'coordinate_cnn', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [7, 16, 16], 'output_shape': [1, 16, 16], 'output_kind': 'diffusion', 'vocab_size': 1, 'noise_steps': 32, 'context_length': 7}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SpatialModel(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, kind: str, condition_dim: int = 0,
                 residual: bool | None = None) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        self.condition_dim = condition_dim
        self.residual = kind in ('residual_cnn', 'spectral_conv') if residual is None else residual
        self.in_proj = nn.Conv2d(input_shape[0] + condition_dim, width, 3, padding=1)
        if kind == 'spectral_conv':
            self.layers = nn.ModuleList(SpectralLayer(width, 6) for _ in range(depth))
        else:
            self.layers = nn.ModuleList(nn.Conv2d(width, width, 3, padding=1) for _ in range(depth))
        classification = len(output_shape) == 1
        self.classification = classification
        self.head = nn.Linear(width, output_shape[0]) if classification else nn.Conv2d(width, output_shape[0], 1)

    def forward(self, x: torch.Tensor, condition: torch.Tensor | None = None,
                timestep: torch.Tensor | None = None, **_: Any) -> torch.Tensor:
        if self.condition_dim:
            if condition is None:
                raise ValueError('condition is required for this denoiser')
            cond = condition.float()
            if timestep is not None:
                cond = torch.cat((cond, timestep.reshape(-1, 1).float()), 1)
            if cond.shape[1] != self.condition_dim:
                raise ValueError('condition feature count does not match candidate spec')
            x = torch.cat((x, cond[:, :, None, None].expand(-1, -1, *x.shape[-2:])), 1)
        h = F.gelu(self.in_proj(x.float()))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        if self.classification:
            return self.head(h.mean((-2, -1)))
        return self.head(h)


class CoordinateCNN(nn.Module):
    def __init__(self, input_shape, output_shape, width, depth, residual):
        super().__init__()
        shape = (input_shape[0] + 2, *input_shape[1:])
        self.net = SpatialModel(shape, output_shape, width, depth, 'cnn', residual=residual)

    def forward(self, x):
        height, width = x.shape[-2:]
        y, horizontal = torch.meshgrid(torch.linspace(-1, 1, height, device=x.device),
                                        torch.linspace(-1, 1, width, device=x.device), indexing='ij')
        coordinates = torch.stack((horizontal, y))[None].expand(x.shape[0], -1, -1, -1)
        return self.net(torch.cat((x, coordinates), dim=1))


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.001
- Weight decay: 0.0

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: diffusion_mse

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: clean_image_denoiser
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [7, 16, 16]
- output_shape: [1, 16, 16]
- output_kind: diffusion
- vocab_size: 1
- noise_steps: 32
- context_length: 7

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'clean_image_denoiser', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [7, 16, 16], 'output_shape': [1, 16, 16], 'output_kind': 'diffusion', 'vocab_size': 1, 'noise_steps': 32, 'context_length': 7}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class CleanImageDenoiser(nn.Module):
    def __init__(self, input_shape, output_shape, width, depth, noise_steps):
        super().__init__()
        channels, height, horizontal = input_shape
        self.noise_steps = noise_steps
        self.condition_dim = channels - 2
        layers = [nn.Conv2d(6, width, 1), nn.GELU()]
        for _ in range(depth):
            layers.extend((nn.Conv2d(width, width, 1), nn.GELU()))
        layers.extend((nn.Conv2d(width, output_shape[0], 1), nn.Sigmoid()))
        self.clean_predictor = nn.Sequential(*layers)

    def forward(self, x):
        batch, _, height, width = x.shape
        y, horizontal = torch.meshgrid(torch.arange(height, device=x.device),
                                        torch.arange(width, device=x.device), indexing='ij')
        condition = x[:, 1:-1]
        radius = condition[:, 3].clamp_min(.1)
        dx = (horizontal[None] - condition[:, 1]) / radius
        dy = (y[None] - condition[:, 2]) / radius
        features = torch.stack((condition[:, 0], dx, dy, dx.square() + dy.square(),
                                torch.maximum(dx.abs(), dy.abs()), condition[:, 4]), dim=1)
        clean = self.clean_predictor(features)
        index = (x[:, -1, 0, 0] * (self.noise_steps - 1)).round().long()
        times = torch.linspace(0, 1, self.noise_steps + 1, device=x.device)
        cumulative = torch.cos((times + .008) / 1.008 * torch.pi / 2).square()
        cumulative = cumulative / cumulative[0]
        beta = (1 - cumulative[1:] / cumulative[:-1]).clamp(1e-5, .999)
        alpha = (1 - beta).cumprod(0)[index, None, None, None]
        return (x[:, :1] - alpha.sqrt() * clean) / (1 - alpha).sqrt()


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.001
- Weight decay: 0.0

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: diffusion_mse

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Direct multi-frame forecasting of moving geometric objects with reflected boundaries. Models receive past frames and jointly predict future frames. This is a direct forecast, not a recursively evaluated learned simulator. Rank by held-out pixel MSE.

### Synthesis (PyTorch)
```python
from __future__ import annotations
import math
import hashlib
from typing import Any
import torch

def synthesize_task(name, params, seed, count):
    cfg = dict(params)
    cfg.setdefault('sensor_seed', int(cfg.get('structure_seed', 20261010)))
    batch = generate('moving_shapes', count, seed, cfg)
    validate('moving_shapes', batch, cfg)
    return batch

def render_shape(conditions: torch.Tensor, size: int=32) -> torch.Tensor:
    """Render [circle=0/square=1, center_x, center_y, radius, intensity]."""
    p = torch.arange(size, dtype=torch.float32)
    (yy, xx) = torch.meshgrid(p, p, indexing='ij')
    c = conditions[:, :, None, None]
    (dx, dy) = (xx - c[:, 1], yy - c[:, 2])
    circle = dx.square() + dy.square() <= c[:, 3].square()
    square = torch.maximum(dx.abs(), dy.abs()) <= c[:, 3]
    mask = torch.where(c[:, 0].bool(), square, circle)
    return (mask * c[:, 4]).float().unsqueeze(1)

def generate(task_id: str, n: int, seed: int, config: dict[str, Any] | None=None) -> dict[str, torch.Tensor]:
    """Generate whole independent examples. Caller owns split seed streams."""
    cfg = config or {}
    if n < 1:
        raise ValueError('n must be positive')
    rng = torch.Generator().manual_seed(seed)
    size = int(cfg.get('size', 32))
    if size < 16:
        raise ValueError('spatial grid must contain at least 16 points per side')
    rng = torch.Generator().manual_seed(seed)
    (context, horizon) = (int(cfg.get('context', 4)), int(cfg.get('horizon', 8)))
    objects = int(cfg.get('objects', 2))
    if context < 2 or horizon < 1 or objects < 1:
        raise ValueError('Motion requires at least two context frames, one future frame and one object')
    radius = size * (0.05 + torch.rand(n, objects, generator=rng) * 0.035)
    position = radius[:, :, None] + torch.rand(n, objects, 2, generator=rng) * (size - 1 - 2 * radius[:, :, None])
    velocity = (torch.rand(n, objects, 2, generator=rng) * 2 - 1) * float(cfg.get('max_speed', 2.0))
    frames = []
    for t in range(context + horizon):
        lower = radius[:, :, None]
        width = size - 1 - 2 * lower
        unfolded = torch.remainder(position + t * velocity - lower, 2 * width)
        centers = lower + width - (unfolded - width).abs()
        image = torch.zeros(n, 1, size, size)
        for j in range(objects):
            cond = torch.stack((torch.full((n,), float(j % 2)), centers[:, j, 0], centers[:, j, 1], radius[:, j], torch.ones(n)), 1)
            image = torch.maximum(image, render_shape(cond, size))
        frames.append(image[:, 0])
    frames = torch.stack(frames, 1)
    return {'x': frames[:, :context], 'y': frames[:, context:], 'position': position, 'velocity': velocity, 'radius': radius}

def validate(task_id: str, batch: dict[str, torch.Tensor], config: dict[str, Any] | None=None) -> dict[str, float | bool]:
    cfg = config or {}
    (x, y) = (batch['x'], batch['y'])
    if not torch.isfinite(x).all() or not torch.isfinite(y).all():
        raise ValueError('Nonfinite spatial dataset')
    if x.shape[0] != y.shape[0] or x.shape[0] < 1:
        raise ValueError('Mismatched spatial example count')
    evidence: dict[str, float | bool] = {'finite': True, 'n': float(x.shape[0])}
    if ((x < 0) | (x > 1)).any() or ((y < 0) | (y > 1)).any():
        raise ValueError('Motion image range outside [0,1]')
    trajectory = torch.cat((x, y), 1)
    (position, velocity, radius) = (batch['position'], batch['velocity'], batch['radius'])
    size = x.shape[-1]
    (yy, xx) = torch.meshgrid(torch.arange(size), torch.arange(size), indexing='ij')
    for t in range(trajectory.shape[1]):
        width = size - 1 - 2 * radius
        centers = radius[..., None] + width[..., None] - (torch.remainder(position + t * velocity - radius[..., None], 2 * width[..., None]) - width[..., None]).abs()
        expected = torch.zeros_like(trajectory[:, t])
        for j in range(radius.shape[1]):
            dx = xx[None] - centers[:, j, 0, None, None]
            dy = yy[None] - centers[:, j, 1, None, None]
            mask = dx.square() + dy.square() <= radius[:, j, None, None].square() if j % 2 == 0 else torch.maximum(dx.abs(), dy.abs()) <= radius[:, j, None, None]
            expected = torch.maximum(expected, mask.float())
        if not torch.equal(trajectory[:, t], expected):
            raise ValueError('Motion target disagrees with reflected trajectory geometry')
    evidence['trajectory_geometry_exact'] = True
    evidence['future_change_fraction'] = float((y[:, 0] != x[:, -1]).float().mean())
    return evidence
TASK = 'moving_shapes'
PARAMS = {'size': 16, 'context': 4, 'horizon': 2, 'objects': 1, 'train_size': 1024, 'test_size': 128, 'validation_size': 64, 'ood_size': 64, 'max_speed': 1, 'instance_seed': 304, 'structure_seed': 10304, 'teacher_seed': 10304, 'ood_shift': {'max_speed': 1.5}, 'input_shape': [4, 16, 16], 'output_shape': [2, 16, 16], 'model_modes': ['mlp', 'cnn', 'residual_cnn', 'fno'], 'objective': 'regression', 'output_kind': 'regression'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_task(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 1024 / 64 / 128
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: regression; model input/output shapes: [4, 16, 16] → [2, 16, 16].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 1000
- batch_size: 16
- total_samples_seen: 16000 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_mse** on the held-out test split.
- Each choice is trained independently for **3** seeds (`0`..`2`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_mse across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: cnn
- width: 32
- depth: 3
- residual: False
- max_parameters: 2000000
- input_shape: [4, 16, 16]
- output_shape: [2, 16, 16]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 4

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'cnn', 'width': 32, 'depth': 3, 'residual': False, 'max_parameters': 2000000, 'input_shape': [4, 16, 16], 'output_shape': [2, 16, 16], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 4}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SpatialModel(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, kind: str, condition_dim: int = 0,
                 residual: bool | None = None) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        self.condition_dim = condition_dim
        self.residual = kind in ('residual_cnn', 'spectral_conv') if residual is None else residual
        self.in_proj = nn.Conv2d(input_shape[0] + condition_dim, width, 3, padding=1)
        if kind == 'spectral_conv':
            self.layers = nn.ModuleList(SpectralLayer(width, 6) for _ in range(depth))
        else:
            self.layers = nn.ModuleList(nn.Conv2d(width, width, 3, padding=1) for _ in range(depth))
        classification = len(output_shape) == 1
        self.classification = classification
        self.head = nn.Linear(width, output_shape[0]) if classification else nn.Conv2d(width, output_shape[0], 1)

    def forward(self, x: torch.Tensor, condition: torch.Tensor | None = None,
                timestep: torch.Tensor | None = None, **_: Any) -> torch.Tensor:
        if self.condition_dim:
            if condition is None:
                raise ValueError('condition is required for this denoiser')
            cond = condition.float()
            if timestep is not None:
                cond = torch.cat((cond, timestep.reshape(-1, 1).float()), 1)
            if cond.shape[1] != self.condition_dim:
                raise ValueError('condition feature count does not match candidate spec')
            x = torch.cat((x, cond[:, :, None, None].expand(-1, -1, *x.shape[-2:])), 1)
        h = F.gelu(self.in_proj(x.float()))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        if self.classification:
            return self.head(h.mean((-2, -1)))
        return self.head(h)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: residual_cnn
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [4, 16, 16]
- output_shape: [2, 16, 16]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 4

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'residual_cnn', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [4, 16, 16], 'output_shape': [2, 16, 16], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 4}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SpatialModel(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, kind: str, condition_dim: int = 0,
                 residual: bool | None = None) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        self.condition_dim = condition_dim
        self.residual = kind in ('residual_cnn', 'spectral_conv') if residual is None else residual
        self.in_proj = nn.Conv2d(input_shape[0] + condition_dim, width, 3, padding=1)
        if kind == 'spectral_conv':
            self.layers = nn.ModuleList(SpectralLayer(width, 6) for _ in range(depth))
        else:
            self.layers = nn.ModuleList(nn.Conv2d(width, width, 3, padding=1) for _ in range(depth))
        classification = len(output_shape) == 1
        self.classification = classification
        self.head = nn.Linear(width, output_shape[0]) if classification else nn.Conv2d(width, output_shape[0], 1)

    def forward(self, x: torch.Tensor, condition: torch.Tensor | None = None,
                timestep: torch.Tensor | None = None, **_: Any) -> torch.Tensor:
        if self.condition_dim:
            if condition is None:
                raise ValueError('condition is required for this denoiser')
            cond = condition.float()
            if timestep is not None:
                cond = torch.cat((cond, timestep.reshape(-1, 1).float()), 1)
            if cond.shape[1] != self.condition_dim:
                raise ValueError('condition feature count does not match candidate spec')
            x = torch.cat((x, cond[:, :, None, None].expand(-1, -1, *x.shape[-2:])), 1)
        h = F.gelu(self.in_proj(x.float()))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        if self.classification:
            return self.head(h.mean((-2, -1)))
        return self.head(h)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: fno
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [4, 16, 16]
- output_shape: [2, 16, 16]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 4

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'fno', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [4, 16, 16], 'output_shape': [2, 16, 16], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 4}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class SpectralLayer(nn.Module):
    def __init__(self, width: int, modes: int) -> None:
        super().__init__()
        self.modes = modes
        scale = 1 / width
        self.weight = nn.Parameter(torch.randn(2, width, width, modes, modes, 2) * scale)
        self.local = nn.Conv2d(width, width, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = torch.fft.rfft2(x)
        out = torch.zeros_like(z)
        m = min(self.modes, z.shape[-2] // 2, z.shape[-1])
        weight = torch.view_as_complex(self.weight.contiguous())
        out[:, :, :m, :m] = torch.einsum('bihw,iohw->bohw', z[:, :, :m, :m], weight[0, :, :, :m, :m])
        out[:, :, -m:, :m] = torch.einsum('bihw,iohw->bohw', z[:, :, -m:, :m], weight[1, :, :, :m, :m])
        return self.local(x) + torch.fft.irfft2(out, s=x.shape[-2:])


class SpatialModel(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, kind: str, condition_dim: int = 0,
                 residual: bool | None = None) -> None:
        super().__init__()
        self.output_shape, self.kind = output_shape, kind
        self.condition_dim = condition_dim
        self.residual = kind in ('residual_cnn', 'spectral_conv') if residual is None else residual
        self.in_proj = nn.Conv2d(input_shape[0] + condition_dim, width, 3, padding=1)
        if kind == 'spectral_conv':
            self.layers = nn.ModuleList(SpectralLayer(width, 6) for _ in range(depth))
        else:
            self.layers = nn.ModuleList(nn.Conv2d(width, width, 3, padding=1) for _ in range(depth))
        classification = len(output_shape) == 1
        self.classification = classification
        self.head = nn.Linear(width, output_shape[0]) if classification else nn.Conv2d(width, output_shape[0], 1)

    def forward(self, x: torch.Tensor, condition: torch.Tensor | None = None,
                timestep: torch.Tensor | None = None, **_: Any) -> torch.Tensor:
        if self.condition_dim:
            if condition is None:
                raise ValueError('condition is required for this denoiser')
            cond = condition.float()
            if timestep is not None:
                cond = torch.cat((cond, timestep.reshape(-1, 1).float()), 1)
            if cond.shape[1] != self.condition_dim:
                raise ValueError('condition feature count does not match candidate spec')
            x = torch.cat((x, cond[:, :, None, None].expand(-1, -1, *x.shape[-2:])), 1)
        h = F.gelu(self.in_proj(x.float()))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        if self.classification:
            return self.head(h.mean((-2, -1)))
        return self.head(h)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Graph shortest-path prediction. Each input packs adjacency followed by a source-node marker; targets are BFS distances from the source. Rank by rounded distance error. An independent Bellman relaxation verifier checks generated targets.

### Synthesis (PyTorch)
```python
from __future__ import annotations
import math
import hashlib
from typing import Any
import torch

def synthesize_task(name, params, seed, count):
    cfg = dict(params)
    cfg.setdefault('sensor_seed', int(cfg.get('structure_seed', 20261010)))
    batch = generate('graph_shortest_path', count, seed, cfg)
    validate('graph_shortest_path', batch, cfg)
    return batch

def generate(task_id: str, n: int, seed: int, config: dict[str, Any] | None=None) -> dict[str, torch.Tensor]:
    """Generate whole independent examples. Caller owns split seed streams."""
    cfg = config or {}
    if n < 1:
        raise ValueError('n must be positive')
    rng = torch.Generator().manual_seed(seed)
    nodes = int(cfg.get('nodes', 32))
    degree = float(cfg.get('mean_degree', 4))
    if nodes < 2 or not 1 <= degree < nodes:
        raise ValueError('Graph needs at least two nodes and degree in [1,nodes)')
    probability = max(0, (degree - 2 * (nodes - 1) / nodes) / (nodes - 1))
    upper = torch.triu((torch.rand(n, nodes, nodes, generator=rng) < probability).float(), 1)
    adjacency = upper + upper.transpose(-1, -2)
    for sample in range(n):
        order = torch.randperm(nodes, generator=rng)
        for j in range(1, nodes):
            parent = int(torch.randint(j, (1,), generator=rng))
            adjacency[sample, order[j], order[parent]] = 1
            adjacency[sample, order[parent], order[j]] = 1
    source = torch.randint(nodes, (n,), generator=rng)
    marker = torch.nn.functional.one_hot(source, nodes).float()
    distances = torch.full((n, nodes), -1.0)
    for sample in range(n):
        queue = [int(source[sample])]
        distances[sample, queue[0]] = 0
        for node in queue:
            for neighbor in adjacency[sample, node].nonzero().flatten().tolist():
                if distances[sample, neighbor] < 0:
                    distances[sample, neighbor] = distances[sample, node] + 1
                    queue.append(neighbor)
    return {'x': torch.cat((adjacency, marker[..., None]), -1), 'y': distances[..., None]}

def validate(task_id: str, batch: dict[str, torch.Tensor], config: dict[str, Any] | None=None) -> dict[str, float | bool]:
    cfg = config or {}
    (x, y) = (batch['x'], batch['y'])
    if not torch.isfinite(x).all() or not torch.isfinite(y).all():
        raise ValueError('Nonfinite spatial dataset')
    if x.shape[0] != y.shape[0] or x.shape[0] < 1:
        raise ValueError('Mismatched spatial example count')
    evidence: dict[str, float | bool] = {'finite': True, 'n': float(x.shape[0])}
    nodes = x.shape[1]
    (adjacency, initial) = (x[:, :, :nodes], x[:, :, nodes])
    if not torch.equal(adjacency, adjacency.transpose(-1, -2)) or adjacency.diagonal(dim1=-2, dim2=-1).any():
        raise ValueError('Graph is not simple undirected')
    if not ((adjacency == 0) | (adjacency == 1)).all() or (adjacency.sum(-1) == 0).any():
        raise ValueError('Invalid graph adjacency')
    if not ((initial == 0) | (initial == 1)).all() or not (initial.sum(-1) == 1).all():
        raise ValueError('BFS needs exactly one marked source')
    distance = y[:, :, 0]
    expected = torch.full_like(distance, float(nodes + 1))
    expected[initial.bool()] = 0
    for _ in range(nodes - 1):
        neighbors = expected[:, None, :].expand(-1, nodes, -1).masked_fill(adjacency == 0, float(nodes + 1))
        expected = torch.minimum(expected, neighbors.min(-1).values + 1)
    if not torch.equal(distance, expected):
        raise ValueError('BFS labels disagree with independent Bellman distance relaxation')
    evidence['distance_verifier_exact'] = True
    return evidence
TASK = 'graph_shortest_path'
PARAMS = {'nodes': 32, 'mean_degree': 4, 'train_size': 512, 'test_size': 128, 'validation_size': 128, 'ood_size': 128, 'instance_seed': 11, 'structure_seed': 10011, 'teacher_seed': 10011, 'ood_shift': {'mean_degree': 2}, 'input_shape': [32, 33], 'output_shape': [32, 1], 'model_modes': ['mlp', 'gnn', 'graph_transformer'], 'objective': 'regression', 'output_kind': 'regression'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_task(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 512 / 128 / 128
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: regression; model input/output shapes: [32, 33] → [32, 1].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 300
- batch_size: 16
- total_samples_seen: 4800 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_distance_error** on the held-out test split.
- Each choice is trained independently for **3** seeds (`0`..`2`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_distance_error across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: mlp
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [32, 33]
- output_shape: [32, 1]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 32

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'mlp', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [32, 33], 'output_shape': [32, 1], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 32}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class FlatMLP(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, residual: bool = False) -> None:
        super().__init__()
        self.output_shape = output_shape
        self.in_proj = nn.Linear(math.prod(input_shape), width)
        self.layers = nn.ModuleList(nn.Linear(width, width) for _ in range(depth))
        self.out_proj = nn.Linear(width, math.prod(output_shape))
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = F.gelu(self.in_proj(x.float().flatten(1)))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        return self.out_proj(h).reshape(x.shape[0], *self.output_shape)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: gnn
- width: 32
- depth: 2
- residual: True
- max_parameters: 2000000
- input_shape: [32, 33]
- output_shape: [32, 1]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 32

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'gnn', 'width': 32, 'depth': 2, 'residual': True, 'max_parameters': 2000000, 'input_shape': [32, 33], 'output_shape': [32, 1], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 32}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class GraphModel(nn.Module):
    def __init__(self, input_dim: int, output_dim: int, width: int, depth: int,
                 kind: str) -> None:
        super().__init__()
        self.kind = kind
        self.embed = nn.Linear(input_dim, width)
        self.layers = nn.ModuleList(nn.Linear(width * 2, width) for _ in range(depth))
        self.head = nn.Linear(width, output_dim)
        if kind == 'graph_attention':
            self.query = nn.Linear(width, width)
            self.key = nn.Linear(width, width)

    def forward(self, x: torch.Tensor, adjacency: torch.Tensor | None = None,
                **_: Any) -> torch.Tensor:
        if adjacency is None:
            raise ValueError('graph candidate requires adjacency')
        h = F.gelu(self.embed(x.float()))
        a = adjacency.float()
        eye = torch.eye(a.shape[-1], device=a.device)[None]
        a = a + eye
        for layer in self.layers:
            if self.kind == 'graph_attention':
                score = self.query(h) @ self.key(h).transpose(-2, -1) / math.sqrt(h.shape[-1])
                weights = score.masked_fill(a <= 0, -torch.inf).softmax(-1)
            else:
                weights = a / a.sum(-1, keepdim=True).clamp_min(1)
            h = h + F.gelu(layer(torch.cat((h, weights @ h), -1)))
        return self.head(h)


class PackedGraphModel(nn.Module):
    def __init__(self, input_shape, output_shape, width, depth, kind):
        super().__init__()
        self.nodes = input_shape[0]
        features = input_shape[-1] - self.nodes
        if features < 1:
            raise ValueError('Packed graph needs adjacency followed by node features')
        self.graph = GraphModel(features, output_shape[-1], width, depth, kind)

    def forward(self, x):
        return self.graph(x[:, :, self.nodes:], adjacency=x[:, :, :self.nodes])


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: gnn
- width: 16
- depth: 2
- residual: True
- max_parameters: 2000000
- input_shape: [32, 33]
- output_shape: [32, 1]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 32

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'gnn', 'width': 16, 'depth': 2, 'residual': True, 'max_parameters': 2000000, 'input_shape': [32, 33], 'output_shape': [32, 1], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 32}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class GraphModel(nn.Module):
    def __init__(self, input_dim: int, output_dim: int, width: int, depth: int,
                 kind: str) -> None:
        super().__init__()
        self.kind = kind
        self.embed = nn.Linear(input_dim, width)
        self.layers = nn.ModuleList(nn.Linear(width * 2, width) for _ in range(depth))
        self.head = nn.Linear(width, output_dim)
        if kind == 'graph_attention':
            self.query = nn.Linear(width, width)
            self.key = nn.Linear(width, width)

    def forward(self, x: torch.Tensor, adjacency: torch.Tensor | None = None,
                **_: Any) -> torch.Tensor:
        if adjacency is None:
            raise ValueError('graph candidate requires adjacency')
        h = F.gelu(self.embed(x.float()))
        a = adjacency.float()
        eye = torch.eye(a.shape[-1], device=a.device)[None]
        a = a + eye
        for layer in self.layers:
            if self.kind == 'graph_attention':
                score = self.query(h) @ self.key(h).transpose(-2, -1) / math.sqrt(h.shape[-1])
                weights = score.masked_fill(a <= 0, -torch.inf).softmax(-1)
            else:
                weights = a / a.sum(-1, keepdim=True).clamp_min(1)
            h = h + F.gelu(layer(torch.cat((h, weights @ h), -1)))
        return self.head(h)


class PackedGraphModel(nn.Module):
    def __init__(self, input_shape, output_shape, width, depth, kind):
        super().__init__()
        self.nodes = input_shape[0]
        features = input_shape[-1] - self.nodes
        if features < 1:
            raise ValueError('Packed graph needs adjacency followed by node features')
        self.graph = GraphModel(features, output_shape[-1], width, depth, kind)

    def forward(self, x):
        return self.graph(x[:, :, self.nodes:], adjacency=x[:, :, :self.nodes])


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.

You are taking the ArchitectureIQ benchmark.

Each question describes one **dataset instance** and several **choices** (each choice is one **candidate**: model, optimizer, loss, and training budget). All choices train on the same dataset instance. Each choice's training budget (`total_samples_seen`) is stated in the prompt — either shared across all choices or specified per choice in the Sample budget section. Your task is to identify which choice will achieve the **best test metric** on the held-out test set after its stated training budget — using architecture and learning-mechanics reasoning only (no training results are provided).

The reply format is defined once, in the final "Your answer" section.

## Dataset
Direct multi-horizon forecasting of diffusion on a graph. Inputs pack adjacency and initial node values. Targets contain future node values generated by the frozen graph transition equation. Rank by held-out MSE.

### Synthesis (PyTorch)
```python
from __future__ import annotations
import math
import hashlib
from typing import Any
import torch

def synthesize_task(name, params, seed, count):
    cfg = dict(params)
    cfg.setdefault('sensor_seed', int(cfg.get('structure_seed', 20261010)))
    batch = generate('graph_diffusion', count, seed, cfg)
    validate('graph_diffusion', batch, cfg)
    return batch

def generate(task_id: str, n: int, seed: int, config: dict[str, Any] | None=None) -> dict[str, torch.Tensor]:
    """Generate whole independent examples. Caller owns split seed streams."""
    cfg = config or {}
    if n < 1:
        raise ValueError('n must be positive')
    rng = torch.Generator().manual_seed(seed)
    nodes = int(cfg.get('nodes', 32))
    degree = float(cfg.get('mean_degree', 4))
    if nodes < 2 or not 1 <= degree < nodes:
        raise ValueError('Graph needs at least two nodes and degree in [1,nodes)')
    probability = max(0, (degree - 2 * (nodes - 1) / nodes) / (nodes - 1))
    upper = torch.triu((torch.rand(n, nodes, nodes, generator=rng) < probability).float(), 1)
    adjacency = upper + upper.transpose(-1, -2)
    for sample in range(n):
        order = torch.randperm(nodes, generator=rng)
        for j in range(1, nodes):
            parent = int(torch.randint(j, (1,), generator=rng))
            adjacency[sample, order[j], order[parent]] = 1
            adjacency[sample, order[parent], order[j]] = 1
    (horizon, alpha) = (int(cfg.get('horizon', 16)), float(cfg.get('alpha', 0.2)))
    if horizon < 1 or not 0 < alpha <= 1:
        raise ValueError('Graph diffusion needs positive horizon and alpha in (0,1]')
    initial = torch.randn(n, nodes, generator=rng)
    current = initial
    transition = adjacency / adjacency.sum(-1, keepdim=True)
    trajectory = []
    for _ in range(horizon):
        current = (1 - alpha) * current + alpha * torch.bmm(transition, current[..., None]).squeeze(-1)
        trajectory.append(current)
    return {'x': torch.cat((adjacency, initial[..., None]), -1), 'y': torch.stack(trajectory, -1)}

def validate(task_id: str, batch: dict[str, torch.Tensor], config: dict[str, Any] | None=None) -> dict[str, float | bool]:
    cfg = config or {}
    (x, y) = (batch['x'], batch['y'])
    if not torch.isfinite(x).all() or not torch.isfinite(y).all():
        raise ValueError('Nonfinite spatial dataset')
    if x.shape[0] != y.shape[0] or x.shape[0] < 1:
        raise ValueError('Mismatched spatial example count')
    evidence: dict[str, float | bool] = {'finite': True, 'n': float(x.shape[0])}
    nodes = x.shape[1]
    (adjacency, initial) = (x[:, :, :nodes], x[:, :, nodes])
    if not torch.equal(adjacency, adjacency.transpose(-1, -2)) or adjacency.diagonal(dim1=-2, dim2=-1).any():
        raise ValueError('Graph is not simple undirected')
    if not ((adjacency == 0) | (adjacency == 1)).all() or (adjacency.sum(-1) == 0).any():
        raise ValueError('Invalid graph adjacency')
    alpha = float(cfg.get('alpha', 0.2))
    previous = torch.cat((initial[:, :, None], y[:, :, :-1]), -1)
    mixed = torch.einsum('bij,bjt->bit', adjacency, previous) / adjacency.sum(-1, keepdim=True)
    expected = (1 - alpha) * previous + alpha * mixed
    if not torch.allclose(y, expected, atol=1e-05):
        raise ValueError('Graph diffusion disagrees with independent transition equation')
    degree = adjacency.sum(-1)
    conserved = (previous * degree[:, :, None]).sum(1)
    mass_error = ((y * degree[:, :, None]).sum(1) - conserved).abs().max()
    evidence.update(transition_max_error=float((y - expected).abs().max()), degree_weighted_mass_error=float(mass_error))
    return evidence
TASK = 'graph_diffusion'
PARAMS = {'nodes': 32, 'mean_degree': 4, 'horizon': 16, 'alpha': 0.2, 'train_size': 512, 'test_size': 128, 'validation_size': 128, 'ood_size': 128, 'instance_seed': 11, 'structure_seed': 10011, 'teacher_seed': 10011, 'ood_shift': {'mean_degree': 2}, 'input_shape': [32, 33], 'output_shape': [32, 16], 'model_modes': ['mlp', 'gnn', 'graph_transformer'], 'objective': 'regression', 'output_kind': 'regression'}

def synthesize_split(seed, count, overrides=None):
    global PARAMS
    original = PARAMS
    PARAMS = dict(PARAMS, **overrides or {})
    try:
        return synthesize_task(TASK, PARAMS, seed, count)
    finally:
        PARAMS = original

def synthesize():
    tr = synthesize_split(PARAMS['instance_seed'] + 20000, PARAMS['train_size'])
    te = synthesize_split(PARAMS['instance_seed'] + 40000, PARAMS['test_size'])
    return (tr['x'], tr['y'], te['x'], te['y'])
```

### Data splits and training protocol
- Train / validation / test size: 512 / 128 / 128
- Split by complete sequences, trajectories, graphs or initial conditions; no overlapping split windows.
- Objective: regression; model input/output shapes: [32, 33] → [32, 16].
- Curves use validation. Final ranking uses held-out test without selecting checkpoints.
- Execution device: cpu.

## Sample budget (same for all choices)
- training_steps: 300
- batch_size: 16
- total_samples_seen: 4800 (= training_steps × batch_size)

## Evaluation metric
- Ground-truth ranking uses **test_mse** on the held-out test split.
- Each choice is trained independently for **3** seeds (`0`..`2`), one `torch.manual_seed(seed)` per run.
- Execution device: cpu.
- The correct choice has the lowest **mean** test_mse across seeds.

## Choices

### Choice A

**Model (natural language)**
- type: expansion_model
- mode: mlp
- width: 32
- depth: 3
- residual: True
- max_parameters: 2000000
- input_shape: [32, 33]
- output_shape: [32, 16]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 32

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'mlp', 'width': 32, 'depth': 3, 'residual': True, 'max_parameters': 2000000, 'input_shape': [32, 33], 'output_shape': [32, 16], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 32}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class FlatMLP(nn.Module):
    def __init__(self, input_shape: tuple[int, ...], output_shape: tuple[int, ...],
                 width: int, depth: int, residual: bool = False) -> None:
        super().__init__()
        self.output_shape = output_shape
        self.in_proj = nn.Linear(math.prod(input_shape), width)
        self.layers = nn.ModuleList(nn.Linear(width, width) for _ in range(depth))
        self.out_proj = nn.Linear(width, math.prod(output_shape))
        self.residual = residual

    def forward(self, x: torch.Tensor, **_: Any) -> torch.Tensor:
        h = F.gelu(self.in_proj(x.float().flatten(1)))
        for layer in self.layers:
            z = F.gelu(layer(h))
            h = h + z if self.residual else z
        return self.out_proj(h).reshape(x.shape[0], *self.output_shape)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice B

**Model (natural language)**
- type: expansion_model
- mode: graph_transformer
- width: 48
- depth: 1
- residual: True
- max_parameters: 2000000
- input_shape: [32, 33]
- output_shape: [32, 16]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 32

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'graph_transformer', 'width': 48, 'depth': 1, 'residual': True, 'max_parameters': 2000000, 'input_shape': [32, 33], 'output_shape': [32, 16], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 32}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class GraphModel(nn.Module):
    def __init__(self, input_dim: int, output_dim: int, width: int, depth: int,
                 kind: str) -> None:
        super().__init__()
        self.kind = kind
        self.embed = nn.Linear(input_dim, width)
        self.layers = nn.ModuleList(nn.Linear(width * 2, width) for _ in range(depth))
        self.head = nn.Linear(width, output_dim)
        if kind == 'graph_attention':
            self.query = nn.Linear(width, width)
            self.key = nn.Linear(width, width)

    def forward(self, x: torch.Tensor, adjacency: torch.Tensor | None = None,
                **_: Any) -> torch.Tensor:
        if adjacency is None:
            raise ValueError('graph candidate requires adjacency')
        h = F.gelu(self.embed(x.float()))
        a = adjacency.float()
        eye = torch.eye(a.shape[-1], device=a.device)[None]
        a = a + eye
        for layer in self.layers:
            if self.kind == 'graph_attention':
                score = self.query(h) @ self.key(h).transpose(-2, -1) / math.sqrt(h.shape[-1])
                weights = score.masked_fill(a <= 0, -torch.inf).softmax(-1)
            else:
                weights = a / a.sum(-1, keepdim=True).clamp_min(1)
            h = h + F.gelu(layer(torch.cat((h, weights @ h), -1)))
        return self.head(h)


class PackedGraphModel(nn.Module):
    def __init__(self, input_shape, output_shape, width, depth, kind):
        super().__init__()
        self.nodes = input_shape[0]
        features = input_shape[-1] - self.nodes
        if features < 1:
            raise ValueError('Packed graph needs adjacency followed by node features')
        self.graph = GraphModel(features, output_shape[-1], width, depth, kind)

    def forward(self, x):
        return self.graph(x[:, :, self.nodes:], adjacency=x[:, :, :self.nodes])


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

### Choice C

**Model (natural language)**
- type: expansion_model
- mode: gnn
- width: 32
- depth: 2
- residual: True
- max_parameters: 2000000
- input_shape: [32, 33]
- output_shape: [32, 16]
- output_kind: regression
- vocab_size: 1
- noise_steps: 100
- context_length: 32

**Model code**
```python
SPEC = {'type': 'expansion_model', 'mode': 'gnn', 'width': 32, 'depth': 2, 'residual': True, 'max_parameters': 2000000, 'input_shape': [32, 33], 'output_shape': [32, 16], 'output_kind': 'regression', 'vocab_size': 1, 'noise_steps': 100, 'context_length': 32}


def _build_module(spec: dict[str, Any]) -> nn.Module:
    kind = spec['mode']
    ins, outs = tuple(spec['input_shape']), tuple(spec['output_shape'])
    width, depth = int(spec['width']), int(spec['depth'])
    if spec['output_kind'] == 'token':
        vocab = int(spec['vocab_size'])
        if kind in {'bigram', 'ngram'}:
            return NgramLM(vocab, width, 1 if kind == 'bigram' else int(spec.get('order', 4)), depth, bool(spec['residual']))
        if kind in {'gru', 'gru_lm'}:
            return TokenModel(CausalGruLM(vocab, int(spec['context_length']), width, depth, bool(spec['residual'])))
        if kind in {'transformer', 'transformer_lm'}:
            return TokenModel(CausalTransformerLM(vocab, int(spec['context_length']), width, depth, 4, width * 2))
        raise ValueError(f'Unknown token architecture: {kind}')
    if kind == 'mlp':
        return FlatMLP(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'coordinate_cnn':
        return CoordinateCNN(ins, outs, width, depth, bool(spec['residual']))
    if kind == 'clean_image_denoiser':
        return CleanImageDenoiser(ins, outs, width, depth, int(spec['noise_steps']))
    if kind in {'gnn', 'graph_transformer'}:
        return PackedGraphModel(ins, outs, width, depth, 'graph_attention' if kind == 'graph_transformer' else 'message_passing')
    if kind in {'gru', 'tcn', 'transformer'}:
        sequence_kind = ('gru_forecaster' if kind == 'gru' else
                         'tcn_forecaster' if kind == 'tcn' else 'transformer_forecaster')
        return SequenceForecaster(ins, outs, width, sequence_kind, depth)
    if kind in {'cnn', 'residual_cnn', 'fno', 'spectral_conv'}:
        spatial_kind = 'spectral_conv' if kind in {'fno', 'spectral_conv'} else 'residual_cnn' if kind == 'residual_cnn' else 'cnn'
        return SpatialModel(ins, outs, width, depth, spatial_kind, residual=bool(spec['residual']))
    if kind in {'vit', 'patch_transformer', 'field_transformer'}:
        return PatchTransformer(ins, outs, width, depth, int(spec.get('patch_size', 4)))
    raise ValueError(f'Unknown expansion architecture: {kind}')


def build_module(spec: dict[str, Any]) -> nn.Module:
    model = _build_module(spec)
    parameters = sum(p.numel() for p in model.parameters())
    cap = int(spec.get('max_parameters', 2_000_000))
    if parameters > cap:
        raise ValueError(f'Candidate exceeds parameter cap: {parameters} > {cap}')
    return model


class GraphModel(nn.Module):
    def __init__(self, input_dim: int, output_dim: int, width: int, depth: int,
                 kind: str) -> None:
        super().__init__()
        self.kind = kind
        self.embed = nn.Linear(input_dim, width)
        self.layers = nn.ModuleList(nn.Linear(width * 2, width) for _ in range(depth))
        self.head = nn.Linear(width, output_dim)
        if kind == 'graph_attention':
            self.query = nn.Linear(width, width)
            self.key = nn.Linear(width, width)

    def forward(self, x: torch.Tensor, adjacency: torch.Tensor | None = None,
                **_: Any) -> torch.Tensor:
        if adjacency is None:
            raise ValueError('graph candidate requires adjacency')
        h = F.gelu(self.embed(x.float()))
        a = adjacency.float()
        eye = torch.eye(a.shape[-1], device=a.device)[None]
        a = a + eye
        for layer in self.layers:
            if self.kind == 'graph_attention':
                score = self.query(h) @ self.key(h).transpose(-2, -1) / math.sqrt(h.shape[-1])
                weights = score.masked_fill(a <= 0, -torch.inf).softmax(-1)
            else:
                weights = a / a.sum(-1, keepdim=True).clamp_min(1)
            h = h + F.gelu(layer(torch.cat((h, weights @ h), -1)))
        return self.head(h)


class PackedGraphModel(nn.Module):
    def __init__(self, input_shape, output_shape, width, depth, kind):
        super().__init__()
        self.nodes = input_shape[0]
        features = input_shape[-1] - self.nodes
        if features < 1:
            raise ValueError('Packed graph needs adjacency followed by node features')
        self.graph = GraphModel(features, output_shape[-1], width, depth, kind)

    def forward(self, x):
        return self.graph(x[:, :, self.nodes:], adjacency=x[:, :, :self.nodes])


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = build_module(SPEC)

    def forward(self, x):
        return self.net(x)
```

**Optimizer**
- Optimizer: Adam
- Learning rate: 0.003
- Weight decay: 0.0
- Betas: [0.9, 0.999]

```python
def build_optimizer(model: nn.Module) -> torch.optim.Optimizer:
    params = model.parameters()
    return torch.optim.Adam(params, lr=0.003, betas=(0.9, 0.999), weight_decay=0.0)
```

**Loss**
- Loss: mean squared error (MSE) on the minibatch

```python
def loss_fn(model: nn.Module, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2)
```

## Your answer
Choose exactly one of A, B, C. End your reply with exactly one tagged field, with nothing after it:

<answer>the letter of your chosen option</answer>

Rules: `<answer>` must contain exactly one letter from {A, B, C}. The grader reads only the **final** `<answer>` tag in your reply.