xai-org/x-algorithm · error · RuntimeError

unexpected head param layout {sorted(params)} (expected {sor

Error message

unexpected head param layout {sorted(params)} (expected {sorted(expected)})

What it means

load_head_params loads head_params.npz, strips each key to its last '/'-segment, and requires exactly the set {w0,b0,w1,b1,w2,b2}. Any other layout (missing layer, extra keys like batch-norm params, different naming) aborts with the sorted key diff.

Source

Thrown at bdsm/runtime/task_heads.py:43

def head_logits(params: dict, x: jax.Array) -> jax.Array:
    n_layers = len(params) // 2
    for i in range(n_layers):
        x = x @ params[f"w{i}"] + params[f"b{i}"]
        if i < n_layers - 1:
            x = jax.nn.gelu(x)
    return x


def load_head_params(head_checkpoint: str) -> dict:
    arrays = np.load(os.path.join(head_checkpoint, "head_params.npz"))
    params = {}
    for key in arrays.files:
        params[key.split("/")[-1]] = arrays[key]
    expected = {"w0", "b0", "w1", "b1", "w2", "b2"}
    if set(params) != expected:
        raise RuntimeError(
            f"unexpected head param layout {sorted(params)} (expected {sorted(expected)})"
        )
    return params

View on GitHub (pinned to 24c60942c5)

Solutions

1. Inspect np.load(...).files of head_params.npz and compare with {w0,b0,w1,b1,w2,b2}
2. Re-export the head with the exact 3-layer wN/bN naming the loader expects
3. If the head architecture legitimately changed, update the expected set in task_heads.py and rebuild
4. Make sure you are pointing at a head checkpoint, not a backbone npz

Example fix

# export heads with exactly these param names:
np.savez('head_params.npz', w0=w0, b0=b0, w1=w1, b1=b1, w2=w2, b2=b2)

arrays = np.load(npz)
names = {k.split('/')[-1] for k in arrays.files}
assert names == {'w0','b0','w1','b1','w2','b2'}, names ^ {'w0','b0','w1','b1','w2','b2'}

def is_expected_head_layout(names: set) -> bool:
    return names == {'w0', 'b0', 'w1', 'b1', 'w2', 'b2'}

try:
    params = load_head_params(ckpt)
except RuntimeError as e:
    if 'unexpected head param layout' in str(e):
        raise SystemExit(f"bad head export: {e}; re-export with wN/bN keys")
    raise

When it happens

Trigger: Loading a head checkpoint exported from a 4-layer head (w3 present), a head with layer-norm keys (gamma/beta), or keys that do not follow the wN/bN naming; npz saved from an older format whose key prefixes collapse to unexpected segments.

Common situations: Architecture change to the head MLP without updating the loader/exporter; loading a backbone param file by mistake; export script renaming keys.

Related errors

• key mismatch vs manifest: missing={missing} extra={extra}
• head registry hash {head_cfg.get('head_registry_hash')} != {
• head checkpoint was trained against a different backbone — r
• head names {unknown} are not in HEAD_ORDER — refusing to sta
• {key}: got {arr.shape}/{arr.dtype}, manifest says {meta['sha

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/175a49bc873e37d9. Report an issue: GitHub.