invoke-ai/InvokeAI · 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's __init__ requires that the positional-embedding axis dims (axes_dim) sum exactly to pe_dim, because EmbedND splits each head's positional embedding across the configured axes. A mismatch means the rotary positional embedding would have the wrong dimensionality, so construction fails.

Source

Thrown at invokeai/backend/flux/model.py:57

    out_channels: Optional[int] = None


class Flux(nn.Module):
    """
    Transformer model for flow matching on sequences.
    """

    def __init__(self, params: FluxParams):
        super().__init__()

        self.params = params
        self.in_channels = params.in_channels
        self.out_channels = params.out_channels or self.in_channels
        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 = nn.Linear(self.in_channels, self.hidden_size, bias=True)
        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
        self.guidance_in = (
            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
        )
        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)

        self.double_blocks = nn.ModuleList(
            [
                DoubleStreamBlock(
                    self.hidden_size,
                    self.num_heads,
                    mlp_ratio=params.mlp_ratio,
                    qkv_bias=params.qkv_bias,

View on GitHub (pinned to 0b6a024f2f)

Solutions

1. Make sum(axes_dim) equal hidden_size // num_heads (standard FLUX: [16, 56, 56] summing to 128)
2. Derive axes_dim programmatically from pe_dim instead of hardcoding
3. Use the predefined params from get_flux_transformers_params(variant)

Example fix

// before
pe_dim = params.hidden_size // params.num_heads  # 128
axes_dim = [128, 128, 128]  # sums to 384 != 128
// after
pe_dim = params.hidden_size // params.num_heads  # 128
axes_dim = [16, 56, 56]  # sums to 128

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

def has_valid_axes_dim(params) -> bool:
    return sum(params.axes_dim) == params.hidden_size // params.num_heads

try:
    model = Flux(params)
except ValueError as e:
    if "expected positional dim" in str(e):
        pe_dim = params.hidden_size // params.num_heads
        params.axes_dim = split_evenly(pe_dim, len(params.axes_dim))
        model = Flux(params)
    else:
        raise

When it happens

Trigger: Constructing the Flux model with FluxParams whose axes_dim list does not sum to hidden_size/num_heads — e.g. axes_dim=[128,128,128] (sum 384) with pe_dim 128, or adding/removing an axis entry for a custom layout without rebalancing.

Common situations: Custom multi-resolution / multi-axis positional encoding experiments; copying axes_dim from FLUX dev (16,56,56 for pe_dim 128) into a model with different hidden_size/num_heads; typos when transcribing params.

Related errors

• Hidden size {params.hidden_size} must be divisible by num_he
• Got {params.axes_dim} but expected positional dim {pe_dim}
• `encoder_hid_dim` has to be defined when `encoder_hid_dim_ty
• User not found or inactive
• Missing authentication credentials

AI-assisted analysis of invoke-ai/InvokeAI@0b6a024f2f (2026-08-29). Data as JSON: /api/errors/2a110e20942e9123. Report an issue: GitHub.