huggingface/pytorch-image-models · error · ValueError

Unsupported patch embedding rank in {checkpoint_path}: {embe

Error message

Unsupported patch embedding rank in {checkpoint_path}: {embed_w.ndim}

What it means

In _load_weights (JAX/Big Vision checkpoint import for ViT), the patch embedding weight must be either a linear (2D) or conv (4D) tensor so it can be converted/resampled to the model's conv patch embed. Any other ndim (e.g. 3D) is unsupported.

Source

Thrown at timm/models/vision_transformer.py:1509

        # Big Vision NaFlex checkpoints use a Dense patch projection over flattened
        # HWC patches, while NaFlexVit may use either Linear or Conv2d projection.
        embeds = model.embeds
        embed_w = _n2p(w[f'{prefix}embedding/kernel'])
        embed_conv_w = None
        if embed_w.ndim == 4:
            embed_conv_w = embed_w
        elif embed_w.ndim == 2:
            # Recover OIHW so input-channel adaptation, patch-size resampling, and
            # the destination patch layout can be handled consistently.
            for checkpoint_in_chans in dict.fromkeys((3, embeds.in_chans)):
                patch_area, remainder = divmod(embed_w.shape[1], checkpoint_in_chans)
                patch_size = int(math.sqrt(patch_area))
                if not remainder and patch_size * patch_size == patch_area:
                    embed_conv_w = embed_w.reshape(
                        embed_w.shape[0], patch_size, patch_size, checkpoint_in_chans).permute(0, 3, 1, 2)
                    break
        else:
            raise ValueError(f'Unsupported patch embedding rank in {checkpoint_path}: {embed_w.ndim}')

        if embed_conv_w is not None:
            embed_conv_w = adapt_input_conv(embeds.in_chans, embed_conv_w)
            if embed_conv_w.shape[-2:] != embeds.patch_size:
                embed_conv_w = resample_patch_embed(
                    embed_conv_w,
                    embeds.patch_size,
                    interpolation=interpolation,
                    antialias=antialias,
                    verbose=True,
                )
            if embeds.is_linear:
                if embeds.channels_last:
                    embed_w = embed_conv_w.permute(0, 2, 3, 1).flatten(1)
                else:
                    embed_w = embed_conv_w.flatten(1)
            else:
                embed_w = embed_conv_w

View on GitHub (pinned to 9a5261e31b)

Solutions

1. Verify the checkpoint is a standard ViT JAX checkpoint with 2D/4D patch embedding weights
2. Pre-convert the checkpoint: reshape the patch kernel to (out, kh, kw, in) or flatten to (out, kh*kw*in)
3. Load the PyTorch (.pth.bin) port of the weights instead of the raw .npz

Example fix

# before
load_checkpoint(model, 'bigvision.npz')  # 3D embedding/kernel
# after
w['embedding/kernel'] = w['embedding/kernel'].reshape(dim, ps, ps, 3)  # pre-reshape
load_checkpoint(model, 'bigvision.npz')

ew = w.get('embedding/kernel')
assert ew is None or ew.ndim in (2, 4), f'unexpected patch embed rank {None if ew is None else ew.ndim}'

try:
    load_pretrained(model, path)
except ValueError as e:
    if 'Unsupported patch embedding rank' in str(e):
        w['embedding/kernel'] = w['embedding/kernel'].reshape(dim, ps, ps, 3)
        load_pretrained(model, path)
    else:
        raise

When it happens

Trigger: Loading a .npz JAX-style ViT checkpoint whose 'embedding/kernel' (or equivalent) has an unexpected rank — e.g. a factorized or pre-reshaped patch tensor.

Common situations: Loading community-converted Big Vision / NaFlex checkpoints where the patch embedding was stored in a non-standard layout.

Related errors

• Patch embedding shape mismatch in {checkpoint_path}: checkpo
• Cannot infer position grid from {pos_embed_w.shape[1]} token
• Unsupported position embedding shape in {checkpoint_path}: {
• Input image must have positive dimensions, got H={height}, W
• Invalid class map file, expected a dict ({class_map_path}).

AI-assisted analysis of huggingface/pytorch-image-models@9a5261e31b (2026-08-27). Data as JSON: /api/errors/ea740ad8d7ea16e8. Report an issue: GitHub.