Comfy-Org/ComfyUI · error · ValueError

Got {params.axes_dim} but expected positional dim {pe_dim}

Error message

Got {params.axes_dim} but expected positional dim {pe_dim}

What it means

After computing pe_dim = hidden_size // num_heads, Flux verifies that the RoPE axes_dim list sums exactly to pe_dim, because each positional axis consumes part of the per-head dimension. A mismatch (sum too large or small) raises ValueError at construction. axes_dim comes from FluxParams and encodes how the head dim is split across (image H, image W, text) axes.

Source

Thrown at comfy/ldm/flux/model.py:82

    """
    Transformer model for flow matching on sequences.
    """

    def __init__(self, image_model=None, final_layer=True, dtype=None, device=None, operations=None, **kwargs):
        super().__init__()
        self.dtype = dtype
        params = FluxParams(**kwargs)
        self.params = params
        self.patch_size = params.patch_size
        self.in_channels = params.in_channels * params.patch_size * params.patch_size
        self.out_channels = params.out_channels * params.patch_size * params.patch_size
        if params.hidden_size % params.num_heads != 0:
            raise ValueError(
                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
            )
        pe_dim = params.hidden_size // params.num_heads
        if sum(params.axes_dim) != pe_dim:
            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
        self.hidden_size = params.hidden_size
        self.num_heads = params.num_heads
        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
        self.img_in = operations.Linear(self.in_channels, self.hidden_size, bias=params.ops_bias, dtype=dtype, device=device)
        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, bias=params.ops_bias, dtype=dtype, device=device, operations=operations)
        if params.vec_in_dim is not None:
            self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size, dtype=dtype, device=device, operations=operations)
        else:
            self.vector_in = None

        self.guidance_in = (
            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, bias=params.ops_bias, dtype=dtype, device=device, operations=operations) if params.guidance_embed else nn.Identity()
        )
        self.txt_in = operations.Linear(params.context_in_dim, self.hidden_size, bias=params.ops_bias, dtype=dtype, device=device)

        if params.txt_norm:
            self.txt_norm = operations.RMSNorm(params.context_in_dim, dtype=dtype, device=device)
        else:

View on GitHub (pinned to 1c6d8d45b3)

Solutions

  1. Make sum(axes_dim) == hidden_size // num_heads; for stock Flux dev/schnell: hidden_size 3072, num_heads 24, axes_dim [16, 56, 56] (sum 128).
  2. When changing hidden_size or num_heads, recompute axes_dim in the same edit — treat the three fields as one unit.
  3. Validate the triple before constructing the model in custom loaders.

Example fix

# before
Flux(hidden_size=1152, num_heads=24, axes_dim=[16, 56, 56], ...)

# after
Flux(hidden_size=1152, num_heads=18, axes_dim=[6, 32, 32], ...)  # 1152/18=64=sum(axes_dim)
Defensive patterns

Strategy: validation

Validate before calling

pe_dim = hidden_size // num_heads
assert sum(axes_dim) == pe_dim, f"sum(axes_dim)={sum(axes_dim)} != hidden_size//num_heads={pe_dim}"

Type guard

def valid_flux_axes(hidden_size: int, num_heads: int, axes_dim: list) -> bool:
    return sum(axes_dim) == hidden_size // num_heads

Prevention

When it happens

Trigger: Overriding hidden_size or num_heads without updating axes_dim (e.g. changing hidden_size 3072->1152 while axes_dim stays [16,56,56] summing to 128); or editing axes_dim for a custom rope layout whose sum no longer equals hidden_size // num_heads.

Common situations: Porting Flux variants (e.g. different resolutions/rope splits) with partially-updated params; config merges that take axes_dim from one variant and hidden_size/num_heads from another.

Related errors


AI-assisted analysis of Comfy-Org/ComfyUI@1c6d8d45b3 (2026-08-14). Data as JSON: /api/errors/2c02190d9210d182. Report an issue: GitHub.