keras-team/keras · error · ValueError

Argument `segment_ids` and `data` should have same leading d

Error message

Argument `segment_ids` and `data` should have same leading dimension. Got {segment_ids_shape} v.s. {data_shape}.

What it means

Segment reduction ops require the first dimension of segment_ids to equal the first dimension of data, since each id labels one row of data. _segment_reduce_validation raises this when both leading dims are statically known (not None) and differ.

Source

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

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}."
        )


class SegmentReduction(Operation):
    def __init__(self, num_segments=None, sorted=False, *, name=None):
        super().__init__(name=name)
        self.num_segments = num_segments
        self.sorted = sorted

    def compute_output_spec(self, data, _):
        output_shape = (self.num_segments,) + tuple(data.shape[1:])
        return KerasTensor(shape=output_shape, dtype=data.dtype)


class SegmentSum(SegmentReduction):

View on GitHub (pinned to 7a34a03db6)

Solutions

1. Regenerate or slice segment_ids from the same filtered data so len(segment_ids) == data.shape[0].
2. Add an explicit check before the call: assert segment_ids.shape[0] == data.shape[0] with a helpful message at your data boundary.
3. If ids derive from labels, verify no rows were added/removed after the label array was materialized (caching pitfalls).

Example fix

// before
data = data[mask]                 # rows filtered
out = ops.segment_sum(data, ids)  # ids still old length -> ValueError

// after
data = data[mask]
ids = ids[mask]                   # keep ids aligned row-for-row
out = ops.segment_sum(data, ids)

from keras import ops

def assert_segment_aligned(data, segment_ids):
    ds, ss = data.shape[0], segment_ids.shape[0]
    assert ds is None or ss is None or ss == ds, (
        f"ids length {ss} != data rows {ds}")

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

def segments_aligned(data, segment_ids) -> bool:
    a, b = data.shape[0], segment_ids.shape[0]
    return a is None or b is None or a == b

When it happens

Trigger: Calling keras.ops.segment_sum(data, ids) where data is (100, 5) but ids has length 80 or 120; slicing data (e.g. data[:100] or data[mask]) without slicing the ids; building ids from a range with an off-by-one against the actual row count.

Common situations: Preprocessing pipelines where filtering drops rows of data but the cached ids array predates the filter; ids computed from a DataFrame after dropping NaNs while data kept all rows; mixing train/val splits between data and ids.

Related errors

• Argument `segment_ids` should be an 1-D vector, got shape: {
• 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/9176b6e2f1c8f5a3. Report an issue: GitHub.