keras-team/keras · error · ValueError

Argument `segment_ids` should be an 1-D vector, got shape: {

Error message

Argument `segment_ids` should be an 1-D vector, got shape: {len(segment_ids_shape)}. Consider either flatten input with segment_ids.reshape((-1)) and data.reshape((-1, ) + data.shape[len(segment_ids.shape):]) or vectorize with vmap.

What it means

The segment reduction ops (keras.ops.segment_sum, segment_max, segment_min, segment_prod) only accept a 1-D segment_ids tensor. _segment_reduce_validation raises this when segment_ids has rank > 1, and the message itself suggests flattening ids and data together or vectorizing with vmap — JAX segment_sum semantics that Keras 3 follows.

Source

Thrown at keras/src/ops/math.py:17

"""Commonly used math operations not included in NumPy."""

from keras.src import backend
from keras.src.api_export import keras_export
from keras.src.backend import KerasTensor
from keras.src.backend import any_symbolic_tensors
from keras.src.backend.common.dtypes import result_type
from keras.src.ops.operation import Operation
from keras.src.ops.operation_utils import broadcast_shapes
from keras.src.ops.operation_utils import reduce_shape


def _segment_reduce_validation(data, segment_ids):
    data_shape = data.shape
    segment_ids_shape = segment_ids.shape
    if len(segment_ids_shape) > 1:
        raise ValueError(
            "Argument `segment_ids` should be an 1-D vector, got shape: "
            f"{len(segment_ids_shape)}. Consider either flatten input with "
            "segment_ids.reshape((-1)) and "
            "data.reshape((-1, ) + data.shape[len(segment_ids.shape):]) or "
            "vectorize with vmap."
        )
    if (
        segment_ids_shape[0] is not None
        and data_shape[0] is not None
        and segment_ids_shape[0] != data_shape[0]
    ):
        raise ValueError(
            "Argument `segment_ids` and `data` should have same leading "
            f"dimension. Got {segment_ids_shape} v.s. "
            f"{data_shape}."
        )

View on GitHub (pinned to 7a34a03db6)

Solutions

1. Flatten as the message suggests: segment_ids = segment_ids.reshape((-1)) and data = data.reshape((-1,) + data.shape[len(segment_ids.shape):]).
2. For per-batch independent reductions, vectorize with keras.ops.vmap (JAX backend) or loop over the batch and call segment_sum per slice.
3. Verify ids are per-element along the flattened leading axis, not a 2-D index grid.

Example fix

// before
from keras import ops
data = ops.ones((4, 10, 3))
ids = ops.tile(ops.arange(10), (4, 1))   # shape (4, 10): 2-D
out = ops.segment_sum(data, ids)         # ValueError

// after
data = ops.ones((4, 10, 3))
ids = ops.tile(ops.arange(10), (4, 1))
ids_flat = ids.reshape((-1,))                      # (40,)
data_flat = data.reshape((-1,) + data.shape[2:])   # (40, 3)
out = ops.segment_sum(data_flat, ids_flat)         # (10, 3)

from keras import ops

def flatten_for_segment(data, segment_ids):
    n_dims = len(segment_ids.shape)
    if n_dims > 1:
        segment_ids = ops.reshape(segment_ids, (-1,))
        data = ops.reshape(data, (-1,) + tuple(data.shape[n_dims:]))
    return data, segment_ids

data, ids = flatten_for_segment(data, ids)
out = ops.segment_sum(data, ids)

import keras

def ids_is_1d(segment_ids) -> bool:
    return keras.ops.ndim(segment_ids) <= 1

When it happens

Trigger: Calling keras.ops.segment_sum(data, segment_ids) with segment_ids of shape (B, N) (one id per element of a batch); using 2-D one-hot or grid-encoded ids; applying grouped reduction over images or time steps while keeping ids multi-dimensional.

Common situations: Porting code from tf.math.unsorted_segment_sum which also demands 1-D ids; grouped pooling in GNN-style layers where node ids arrive as (batch, nodes); assuming extra id dims batch the op when vmap is the intended mechanism.

Related errors

• Argument `segment_ids` and `data` should have same leading d
• The `weights` argument should be either `None` (random initi
• Expected mode to be one of `caffe`, `tf` or `torch`. Receive
• TF-IDF data must be a 1-index array. Received: type(idf_weig
• When using `output_mode={self.output_mode}` and `pad_to_max_

AI-assisted analysis of keras-team/keras@7a34a03db6 (2026-08-25). Data as JSON: /api/errors/86ac81a57ab19fec. Report an issue: GitHub.