xai-org/x-algorithm · error · ValueError

Tensor {name!r}: checkpoint has shape {tuple(t.shape)}, but

Error message

Tensor {name!r}: checkpoint has shape {tuple(t.shape)}, but initialized state has shape {tuple(dest.shape)}. Use 'no_loading' to skip this tensor or 'domains' to load a partial slice.

What it means

_open_tensor opens a tensor from the checkpoint and, when no domain restriction is given, requires its shape to equal the destination (initialized state) tensor's shape. A mismatch raises this ValueError suggesting 'no_loading' or 'domains'.

Source

Thrown at phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py:219

    use_zarr3: bool,
    ts_context: ts.Context,
    dest: jax.Array,
    has_domain: bool,
    tspec_transform: Callable[[dict[str, Any]], dict[str, Any]] | None,
) -> ts.TensorStore:
    info = ocp.type_handlers.ParamInfo(
        name=checkpoint_name,
        path=path / checkpoint_name,
        parent_dir=path,
        is_ocdbt_checkpoint=True,
        use_zarr3=use_zarr3,
    )
    tspec = ocp.type_handlers.get_json_tspec_read(info, use_ocdbt=True)
    if tspec_transform is not None:
        tspec = tspec_transform(tspec)
    t = ts.open(ts.Spec(tspec), open=True, context=ts_context).result()
    if not has_domain and tuple(t.shape) != tuple(dest.shape):
        raise ValueError(
            f"Tensor {name!r}: checkpoint has shape {tuple(t.shape)}, "
            f"but initialized state has shape {tuple(dest.shape)}. "
            f"Use 'no_loading' to skip this tensor or 'domains' to load a partial slice."
        )
    return t


def _read_into_shards(
    t: ts.TensorStore,
    array: jax.Array,
    mask: list[bool],
    restricted_domain: ts.IndexDomain | None = None,
):
    memory_kind = array.sharding.memory_kind
    is_cpu = all(d.platform == "cpu" for d in array.sharding.addressable_devices)
    assert memory_kind == "pinned_host" or is_cpu, (
        f"expected pinned_host memory, got {memory_kind!r}"
    )

View on GitHub (pinned to 24c60942c5)

Solutions

1. Revert the config change so shapes match the checkpoint
2. Add the tensor to no_loading to skip restoring it
3. Use domains to load the overlapping slice if sizes changed intentionally (e.g. extended vocab)
4. Re-save the checkpoint after the architecture change so shapes align

Example fix

# before
load_checkpoint(path, host_state, no_loading=set())
# after
load_checkpoint(path, host_state, no_loading={"embedding"})  # vocab resized

# compare shapes before loading
for name, dest in state.items():
    ckpt_shape = checkpoint_index[name].shape  # from metadata
    if tuple(ckpt_shape) != tuple(dest.shape):
        print(f"shape mismatch: {name} {ckpt_shape} vs {dest.shape}")

try:
    load_checkpoint(...)
except ValueError as e:
    if "checkpoint has shape" in str(e):
        load_checkpoint(..., no_loading=offending_tensors)

When it happens

Trigger: load_checkpoint where the model was initialized with different shapes than saved — e.g. vocab_size, mesh/sharding changes, or num_layers changed in the config while reusing an old checkpoint.

Common situations: Resuming after a config change that alters a tensor's axis size, loading a sharded checkpoint into a differently-shaped mesh, or model architecture version changes.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

• {key}: got {arr.shape}/{arr.dtype}, manifest says {meta['sha
• unexpected head param layout {sorted(params)} (expected {sor
• head registry hash {head_cfg.get('head_registry_hash')} != {
• head checkpoint was trained against a different backbone — r
• key mismatch vs manifest: missing={missing} extra={extra}

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