xai-org/x-algorithm · error · ValueError

Expected 2D emb_table, got shape={embeddings.shape}

Error message

Expected 2D emb_table, got shape={embeddings.shape}

What it means

Raised by _maybe_init_rows during lazy initialization of embedding table rows. The function expects the emb_table to be a 2D array of shape [num_rows, row_size] so it can index embeddings[idx, 0] and fill whole rows; if the array is 1D or 3D+, shape[1] is not the row width and row writes would be wrong, so it fails fast. It surfaces through lookup_h2d_embeddings when the passed emb_table has the wrong rank.

Source

Thrown at phoenix/xrex/inference/h2d.py:396

        candidate_multimodal_embeddings_gpu = mm_2d.reshape(
            total_batch, mm_cand_seq, mm_buf.embedding_dim
        )

    return merged_device, candidate_multimodal_embeddings_gpu


def _maybe_init_rows(
    embeddings: npt.NDArray,
    row_indexes: npt.NDArray[np.uint32],
) -> None:
    logger.info(
        "Lazy init emb_table rows: unique_rows=%d shape=%s dtype=%s",
        len(np.unique(row_indexes)),
        embeddings.shape,
        embeddings.dtype,
    )
    if embeddings.ndim != 2:
        raise ValueError(f"Expected 2D emb_table, got shape={embeddings.shape}")
    row_size = embeddings.shape[1]
    scale = 1.0 / math.sqrt(row_size)
    unique_rows = np.unique(row_indexes)
    for row in unique_rows:
        idx = int(row)
        first = int(embeddings[idx, 0])
        if first != 0:
            continue
        rng = np.random.default_rng(idx)
        values = rng.uniform(-scale, scale, size=row_size).astype(np.float32)
        if embeddings.dtype == np.uint16:
            embeddings[idx] = (values.view(np.uint32) >> 16).astype(np.uint16)
        elif embeddings.dtype == np.uint32:
            embeddings[idx] = values.view(np.uint32)
        else:
            raise ValueError(f"Unsupported emb_table dtype: {embeddings.dtype}")

View on GitHub (pinned to 24c60942c5)

Solutions

  1. Reshape the table to 2D before lookup: embeddings.reshape(num_rows, row_size) with the correct vocab and dim from the config
  2. If you intended per-item embeddings, slice out the actual table (e.g. batch[0]) rather than passing the stacked array
  3. Verify the checkpoint field you loaded is the emb_table itself and not a flattened byte buffer

Example fix

# before
table = np.frombuffer(blob, dtype=np.uint16)  # 1D
rows = lookup_h2d_embeddings(idx, table, ...)
# after
table = np.frombuffer(blob, dtype=np.uint16).reshape(vocab, dim)
rows = lookup_h2d_embeddings(idx, table, ...)
Defensive patterns

Strategy: validation

Validate before calling

assert isinstance(embeddings, np.ndarray) and embeddings.ndim == 2, (
    f"emb_table must be 2D, got {embeddings.shape}")

Type guard

def is_2d_table(a: np.ndarray) -> bool:
    return isinstance(a, np.ndarray) and a.ndim == 2

Try / catch

try:
    rows = lookup_h2d_embeddings(idx, table, ...)
except ValueError as e:
    if "Expected 2D" in str(e):
        table = table.reshape(vocab, dim)
    else:
        raise

Prevention

When it happens

Trigger: Calling lookup_h2d_embeddings with an embeddings array that is 1D (a single flattened row), 3D (e.g. [vocab, seq, dim] not squeezed), or a 0-d array; also when a checkpoint loader reshapes the table incorrectly before lookup.

Common situations: Loading a raw flattened embedding blob from a checkpoint and passing it without reshape; passing per-sequence embeddings (batch of 2D tables) instead of the vocab table; test fixtures constructing np.zeros(vocab*dim) instead of (vocab, dim).

Related errors

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