Comfy-Org/ComfyUI · critical · 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 checking head divisibility, Chroma derives pe_dim = hidden_size // num_heads and requires the rotational-position-embedding axes_dim list to sum exactly to pe_dim: each axis gets a slice of the per-head dim in EmbedND. If sum(axes_dim) != pe_dim, RoPE would be built over the wrong dimensionality, so the constructor rejects the config. This is the same family of check as [121]: config-derived parameters must be internally consistent.

Source

Thrown at comfy/ldm/chroma/model.py:67

    """
    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 = ChromaParams(**kwargs)
        self.params = params
        self.patch_size = params.patch_size
        self.in_channels = params.in_channels
        self.out_channels = params.out_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.in_dim = params.in_dim
        self.out_dim = params.out_dim
        self.hidden_dim = params.hidden_dim
        self.n_layers = params.n_layers
        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=True, dtype=dtype, device=device)
        self.txt_in = operations.Linear(params.context_in_dim, self.hidden_size, dtype=dtype, device=device)
        # set as nn identity for now, will overwrite it later.
        self.distilled_guidance_layer = Approximator(
                    in_dim=self.in_dim,
                    hidden_dim=self.hidden_dim,
                    out_dim=self.out_dim,
                    n_layers=self.n_layers,
                    dtype=dtype, device=device, operations=operations
                )

View on GitHub (pinned to 1c6d8d45b3)

Solutions

1. Recompute axes_dim so its entries sum to hidden_size // num_heads (pe_dim)
2. If editing one axes_dim entry (e.g. the time axis), adjust another entry so the total stays equal to pe_dim
3. Restore the official Chroma axes_dim values from the checkpoint config instead of hand-writing them

Example fix

# before
# hidden_size=3000, num_heads=24 -> pe_dim=125
axes_dim = [16, 56, 56]  # sums to 128, raises

# after
pe_dim = hidden_size // num_heads
t_axes = pe_dim - 56 - 56  # keep image axes at 56 each, derive the rest
axes_dim = [t_axes, 56, 56]

pe_dim = hidden_size // num_heads
if sum(axes_dim) != pe_dim:
    axes_dim = [pe_dim - sum(axes_dim[1:])] + list(axes_dim[1:])  # rebalance first axis
assert sum(axes_dim) == pe_dim

When it happens

Trigger: Constructing Chroma with axes_dim whose entries do not sum to hidden_size // num_heads. E.g. hidden_size=3000, num_heads=24 gives pe_dim=125 but axes_dim=[16,56,56] sums to 128 - the constructor raises.

Common situations: Porting a config from a different DiT (Flux-style axes_dim=[16,56,56] with 128 head dim) into Chroma; editing axes_dim when customizing positional encoding without recomputing pe_dim; loading a misidentified checkpoint.

Related errors

• Hidden size {params.hidden_size} must be divisible by num_he
• Got {params.axes_dim} but expected positional dim {pe_dim}
• Block type {block_type} not supported
• Hidden size {params.hidden_size} must be divisible by num_he
• Unsupported nerf_final_head_type {params.nerf_final_head_typ

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