run-llama/llama_index · error · NotImplementedError

Unhandled shape {array.shape}.

Error message

Unhandled shape {array.shape}.

What it means

PoolEmbedding.cls_pooling reduces an embedding tensor to its [CLS] token representation and only understands 3-D (batch, tokens, dim) and 2-D (tokens, dim) arrays. A 1-D vector (single already-pooled embedding) or 4-D+ array has no token axis to select from, so pooling is not implemented for it and NotImplementedError is raised.

Source

Thrown at llama-index-core/llama_index/core/embeddings/pooling.py:40

    @overload
    def cls_pooling(cls, array: np.ndarray) -> np.ndarray: ...

    @classmethod
    @overload
    # TODO: Remove this `type: ignore` after the false positive problem
    #  is addressed in mypy: https://github.com/python/mypy/issues/15683 .
    def cls_pooling(cls, array: "torch.Tensor") -> "torch.Tensor":  # type: ignore
        ...

    @classmethod
    def cls_pooling(
        cls, array: "Union[np.ndarray, torch.Tensor]"
    ) -> "Union[np.ndarray, torch.Tensor]":
        if len(array.shape) == 3:
            return array[:, 0]
        if len(array.shape) == 2:
            return array[0]
        raise NotImplementedError(f"Unhandled shape {array.shape}.")

    @classmethod
    def mean_pooling(cls, array: np.ndarray) -> np.ndarray:
        if len(array.shape) == 3:
            return array.mean(axis=1)
        if len(array.shape) == 2:
            return array.mean(axis=0)
        raise NotImplementedError(f"Unhandled shape {array.shape}.")

View on GitHub (pinned to afd0fef371)

Solutions

  1. Ensure input is (batch, tokens, dim) — keep the batch dimension even for one text: array[None, ...] or tensor.unsqueeze(0)
  2. If vectors are already pooled (1-D), skip pooling entirely and return them as-is
  3. Reshape 4-D inputs to 3-D by merging leading dimensions before pooling

Example fix

// before
pooled = PoolEmbedding.cls_pooling(vec)  # vec.shape == (384,) -> NotImplementedError

// after
if vec.ndim == 1:
    pooled = vec  # already pooled
else:
    pooled = PoolEmbedding.cls_pooling(vec[None, :] if vec.ndim == 2 else vec)
Defensive patterns

Strategy: type-guard

Validate before calling

assert array.ndim in (2, 3), f"expected 2D/3D, got {array.shape}"

Type guard

def is_poolable(array) -> bool:
    return getattr(array, "ndim", 0) in (2, 3)

Try / catch

try:
    pooled = PoolEmbedding.cls_pooling(array)
except NotImplementedError:
    pooled = array  # assume already-pooled 1-D vector

Prevention

When it happens

Trigger: Passing a single 1-D embedding vector (shape (dim,)) into cls_pooling; feeding 4-D image-feature tensors; calling get_text_embedding on a model whose backend already returns pooled 1-D vectors.

Common situations: Writing a custom multi-modal embedding around PoolEmbedding; backends/version changes that return pre-pooled vectors; batching logic that squeezes the wrong axis.

Related errors


AI-assisted analysis of run-llama/llama_index@afd0fef371 (2026-08-15). Data as JSON: /api/errors/ad7319b13b9e1cdf. Report an issue: GitHub.