xai-org/x-algorithm · error · ValueError

{emb_size=} was divided into {output_vocab_size=} equal part

Error message

{emb_size=} was divided into {output_vocab_size=} equal parts of length {D_per_V} each. But received an input with {V=}.

What it means

multi_hot_to_embeddings splits the embedding table row of size emb_size into output_vocab_size equal chunks (D_per_V each) and requires the incoming multi-hot input's vocab dimension V to satisfy V * D_per_V == D (the emb_size). A mismatch means the input vocabulary size differs from the one the table was reshaped for, and the error reports all four numbers.

Source

Thrown at phoenix/xrex/models/recsys_model.py:1909

        embedding_table = get_parameter(
            name,
            shape=[
                1,
                emb_size,
            ],
            init=embed_init,
            dtype=jnp.float32,
            pspec=P(),
            rms_clip_axes=(-2, -1),
        )
        table_reshaped = embedding_table.reshape(output_vocab_size, emb_size // output_vocab_size)

        B, S, V = input.shape
        D = emb_size
        D_per_V = table_reshaped.shape[1]

        if V * D_per_V != D:
            raise ValueError(
                f"{emb_size=} was divided into {output_vocab_size=} equal parts of length {D_per_V} each. But received an input with {V=}."
            )

        input_reshaped = (2 * input - 1)[:, :, :, None]
        table_reshaped = table_reshaped[None, None, :, :]

        selected_embeddings = input_reshaped * table_reshaped

        output = selected_embeddings.reshape(B, S, D)

        mask = jnp.any(input, axis=-1)
        output = output * mask[..., None]
        output = output.astype(DTYPE_BY_NAME[self.config.fprop_dtype])

        return output, embedding_table

    @hk.transparent
    def single_hot_to_embeddings(

View on GitHub (pinned to 24c60942c5)

Solutions

1. Align V in the input batch with output_vocab_size (re-tokenize / re-vocab the data).
2. Update output_vocab_size in the config to the actual input V.
3. Ensure emb_size is divisible by output_vocab_size and rebuild the table if either changed.

Example fix

# before
output_vocab_size = 1000   # input has V=1200

# after
output_vocab_size = 1200     # matches input.shape[-1]

B, S, V = multi_hot.shape
assert V == output_vocab_size and emb_size % output_vocab_size == 0, (V, output_vocab_size, emb_size)

When it happens

Trigger: Feeding a [B, S, V] multi-hot tensor whose V differs from output_vocab_size used when the table was reshaped; changing the vocab size in data preprocessing without regenerating the embedding table; emb_size not evenly divisible by output_vocab_size plus a stale V.

Common situations: Vocabulary grown/shrunk between training and serving; using an old checkpoint's embeddings with a new tokenizer/vocab; mismatched output_vocab_size config vs. input pipeline.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

• Only 1D arrays are supported for unique.
• Block sparse tensors{context} must have shapes (B, H, M) and
• Block sparse tensors{context} {dim_name} dim must be {tgt} o
• Block sparse tensors{context} must share the same m-block di
• Not a canonical UMMA_K Layout: Expected MN-size multiple of

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/d6556f5e7b68050c. Report an issue: GitHub.