{"record":{"id":"ea740ad8d7ea16e8","repo":"huggingface/pytorch-image-models","slug":"unsupported-patch-embedding-rank-in-checkpoint-pa","errorCode":null,"errorMessage":"Unsupported patch embedding rank in {checkpoint_path}: {embed_w.ndim}","messagePattern":"Unsupported patch embedding rank in (.+?): (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"timm/models/vision_transformer.py","lineNumber":1509,"sourceCode":"        # Big Vision NaFlex checkpoints use a Dense patch projection over flattened\n        # HWC patches, while NaFlexVit may use either Linear or Conv2d projection.\n        embeds = model.embeds\n        embed_w = _n2p(w[f'{prefix}embedding/kernel'])\n        embed_conv_w = None\n        if embed_w.ndim == 4:\n            embed_conv_w = embed_w\n        elif embed_w.ndim == 2:\n            # Recover OIHW so input-channel adaptation, patch-size resampling, and\n            # the destination patch layout can be handled consistently.\n            for checkpoint_in_chans in dict.fromkeys((3, embeds.in_chans)):\n                patch_area, remainder = divmod(embed_w.shape[1], checkpoint_in_chans)\n                patch_size = int(math.sqrt(patch_area))\n                if not remainder and patch_size * patch_size == patch_area:\n                    embed_conv_w = embed_w.reshape(\n                        embed_w.shape[0], patch_size, patch_size, checkpoint_in_chans).permute(0, 3, 1, 2)\n                    break\n        else:\n            raise ValueError(f'Unsupported patch embedding rank in {checkpoint_path}: {embed_w.ndim}')\n\n        if embed_conv_w is not None:\n            embed_conv_w = adapt_input_conv(embeds.in_chans, embed_conv_w)\n            if embed_conv_w.shape[-2:] != embeds.patch_size:\n                embed_conv_w = resample_patch_embed(\n                    embed_conv_w,\n                    embeds.patch_size,\n                    interpolation=interpolation,\n                    antialias=antialias,\n                    verbose=True,\n                )\n            if embeds.is_linear:\n                if embeds.channels_last:\n                    embed_w = embed_conv_w.permute(0, 2, 3, 1).flatten(1)\n                else:\n                    embed_w = embed_conv_w.flatten(1)\n            else:\n                embed_w = embed_conv_w","sourceCodeStart":1491,"sourceCodeEnd":1527,"githubUrl":"https://github.com/huggingface/pytorch-image-models/blob/9a5261e31b3b5128526eb2658333b4c0a54464ae/timm/models/vision_transformer.py#L1491-L1527","documentation":"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.","triggerScenarios":"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.","commonSituations":"Loading community-converted Big Vision / NaFlex checkpoints where the patch embedding was stored in a non-standard layout.","solutions":["Verify the checkpoint is a standard ViT JAX checkpoint with 2D/4D patch embedding weights","Pre-convert the checkpoint: reshape the patch kernel to (out, kh, kw, in) or flatten to (out, kh*kw*in)","Load the PyTorch (.pth.bin) port of the weights instead of the raw .npz"],"exampleFix":"# before\nload_checkpoint(model, 'bigvision.npz')  # 3D embedding/kernel\n# after\nw['embedding/kernel'] = w['embedding/kernel'].reshape(dim, ps, ps, 3)  # pre-reshape\nload_checkpoint(model, 'bigvision.npz')","handlingStrategy":"validation","validationCode":"ew = w.get('embedding/kernel')\nassert ew is None or ew.ndim in (2, 4), f'unexpected patch embed rank {None if ew is None else ew.ndim}'","typeGuard":null,"tryCatchPattern":"try:\n    load_pretrained(model, path)\nexcept ValueError as e:\n    if 'Unsupported patch embedding rank' in str(e):\n        w['embedding/kernel'] = w['embedding/kernel'].reshape(dim, ps, ps, 3)\n        load_pretrained(model, path)\n    else:\n        raise","preventionTips":["Inspect .npz tensor shapes before loading","Prefer official PyTorch ports of JAX checkpoints"],"tags":["timm","vit","checkpoint","jax"],"backgroundTag":"checkpoint-weight-mismatch","analyzedSha":"9a5261e31b3b5128526eb2658333b4c0a54464ae","analyzedAt":"2026-08-27T02:34:25.417Z","schemaVersion":2},"datasetVersion":"2026-08-27T03:17:27.898Z"}