{"record":{"id":"9176b6e2f1c8f5a3","repo":"keras-team/keras","slug":"argument-segment-ids-and-data-should-have-same","errorCode":null,"errorMessage":"Argument `segment_ids` and `data` should have same leading dimension. Got {segment_ids_shape} v.s. {data_shape}.","messagePattern":"Argument `segment_ids` and `data` should have same leading dimension\\. Got (.+?) v\\.s\\. (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"keras/src/ops/math.py","lineNumber":29,"sourceCode":"\n\ndef _segment_reduce_validation(data, segment_ids):\n    data_shape = data.shape\n    segment_ids_shape = segment_ids.shape\n    if len(segment_ids_shape) > 1:\n        raise ValueError(\n            \"Argument `segment_ids` should be an 1-D vector, got shape: \"\n            f\"{len(segment_ids_shape)}. Consider either flatten input with \"\n            \"segment_ids.reshape((-1)) and \"\n            \"data.reshape((-1, ) + data.shape[len(segment_ids.shape):]) or \"\n            \"vectorize with vmap.\"\n        )\n    if (\n        segment_ids_shape[0] is not None\n        and data_shape[0] is not None\n        and segment_ids_shape[0] != data_shape[0]\n    ):\n        raise ValueError(\n            \"Argument `segment_ids` and `data` should have same leading \"\n            f\"dimension. Got {segment_ids_shape} v.s. \"\n            f\"{data_shape}.\"\n        )\n\n\nclass SegmentReduction(Operation):\n    def __init__(self, num_segments=None, sorted=False, *, name=None):\n        super().__init__(name=name)\n        self.num_segments = num_segments\n        self.sorted = sorted\n\n    def compute_output_spec(self, data, _):\n        output_shape = (self.num_segments,) + tuple(data.shape[1:])\n        return KerasTensor(shape=output_shape, dtype=data.dtype)\n\n\nclass SegmentSum(SegmentReduction):","sourceCodeStart":11,"sourceCodeEnd":47,"githubUrl":"https://github.com/keras-team/keras/blob/7a34a03db60bf60042242d6a556fc3be119046a5/keras/src/ops/math.py#L11-L47","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Regenerate or slice segment_ids from the same filtered data so len(segment_ids) == data.shape[0].","Add an explicit check before the call: assert segment_ids.shape[0] == data.shape[0] with a helpful message at your data boundary.","If ids derive from labels, verify no rows were added/removed after the label array was materialized (caching pitfalls)."],"exampleFix":"// before\ndata = data[mask]                 # rows filtered\nout = ops.segment_sum(data, ids)  # ids still old length -> ValueError\n\n// after\ndata = data[mask]\nids = ids[mask]                   # keep ids aligned row-for-row\nout = ops.segment_sum(data, ids)","handlingStrategy":"validation","validationCode":"from keras import ops\n\ndef assert_segment_aligned(data, segment_ids):\n    ds, ss = data.shape[0], segment_ids.shape[0]\n    assert ds is None or ss is None or ss == ds, (\n        f\"ids length {ss} != data rows {ds}\")\n\nassert_segment_aligned(data, ids)\nout = ops.segment_sum(data, ids)","typeGuard":"def segments_aligned(data, segment_ids) -> bool:\n    a, b = data.shape[0], segment_ids.shape[0]\n    return a is None or b is None or a == b","tryCatchPattern":null,"preventionTips":["Apply the same mask/slice to data and ids in one function.","Regenerate ids whenever data is filtered; never cache ids across preprocessing changes.","Property-test len(ids) == data.shape[0] over random subsets."],"tags":["keras","segment-ops","shape-validation","data-alignment","grouped-reduction"],"backgroundTag":"tensor-rank-or-shape-mismatch","analyzedSha":"7a34a03db60bf60042242d6a556fc3be119046a5","analyzedAt":"2026-08-25T21:25:25.994Z","schemaVersion":2},"datasetVersion":"2026-08-26T02:17:13.382Z"}