{"record":{"id":"86ac81a57ab19fec","repo":"keras-team/keras","slug":"argument-segment-ids-should-be-an-1-d-vector-go","errorCode":null,"errorMessage":"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.","messagePattern":"Argument `segment_ids` should be an 1-D vector, got shape: (.+?)\\. Consider either flatten input with segment_ids\\.reshape\\(\\(-1\\)\\) and data\\.reshape\\(\\(-1, \\) \\+ data\\.shape\\[len\\(segment_ids\\.shape\\):\\]\\) or vectorize with vmap\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"keras/src/ops/math.py","lineNumber":17,"sourceCode":"\"\"\"Commonly used math operations not included in NumPy.\"\"\"\n\nfrom keras.src import backend\nfrom keras.src.api_export import keras_export\nfrom keras.src.backend import KerasTensor\nfrom keras.src.backend import any_symbolic_tensors\nfrom keras.src.backend.common.dtypes import result_type\nfrom keras.src.ops.operation import Operation\nfrom keras.src.ops.operation_utils import broadcast_shapes\nfrom keras.src.ops.operation_utils import reduce_shape\n\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","sourceCodeStart":1,"sourceCodeEnd":35,"githubUrl":"https://github.com/keras-team/keras/blob/7a34a03db60bf60042242d6a556fc3be119046a5/keras/src/ops/math.py#L1-L35","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Flatten as the message suggests: segment_ids = segment_ids.reshape((-1)) and data = data.reshape((-1,) + data.shape[len(segment_ids.shape):]).","For per-batch independent reductions, vectorize with keras.ops.vmap (JAX backend) or loop over the batch and call segment_sum per slice.","Verify ids are per-element along the flattened leading axis, not a 2-D index grid."],"exampleFix":"// before\nfrom keras import ops\ndata = ops.ones((4, 10, 3))\nids = ops.tile(ops.arange(10), (4, 1))   # shape (4, 10): 2-D\nout = ops.segment_sum(data, ids)         # ValueError\n\n// after\ndata = ops.ones((4, 10, 3))\nids = ops.tile(ops.arange(10), (4, 1))\nids_flat = ids.reshape((-1,))                      # (40,)\ndata_flat = data.reshape((-1,) + data.shape[2:])   # (40, 3)\nout = ops.segment_sum(data_flat, ids_flat)         # (10, 3)","handlingStrategy":"validation","validationCode":"from keras import ops\n\ndef flatten_for_segment(data, segment_ids):\n    n_dims = len(segment_ids.shape)\n    if n_dims > 1:\n        segment_ids = ops.reshape(segment_ids, (-1,))\n        data = ops.reshape(data, (-1,) + tuple(data.shape[n_dims:]))\n    return data, segment_ids\n\ndata, ids = flatten_for_segment(data, ids)\nout = ops.segment_sum(data, ids)","typeGuard":"import keras\n\ndef ids_is_1d(segment_ids) -> bool:\n    return keras.ops.ndim(segment_ids) <= 1","tryCatchPattern":null,"preventionTips":["Always generate ids with ops.arange/ops.tile ending at rank 1.","Batch via vmap instead of 2-D ids.","Centralize the flatten step in a helper reused by all segment_* calls."],"tags":["keras","segment-ops","shape-validation","grouped-reduction","jax-semantics"],"backgroundTag":"tensor-rank-or-shape-mismatch","analyzedSha":"7a34a03db60bf60042242d6a556fc3be119046a5","analyzedAt":"2026-08-25T21:25:25.994Z","schemaVersion":2},"datasetVersion":"2026-08-26T02:17:13.382Z"}