{"record":{"id":"da6ef01aeef60555","repo":"jax-ml/jax","slug":"invalid-shape-for-q-segment-ids-segment-ids-q-sh","errorCode":null,"errorMessage":"Invalid shape for q segment_ids: {segment_ids.q.shape}. Expected: {(q_seq_len,)}","messagePattern":"Invalid shape for q segment_ids: (.+?)\\. Expected: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":979,"sourceCode":"  if k.shape[:-1] != v.shape[:-1]:\n    raise ValueError(\n        f\"Expected 'key' {k.shape} and 'value' {v.shape} to have the same \"\n        \"leading dimensions.\"\n    )\n\n  assert bkv_compute is not None\n  if bkv % bkv_compute:\n    raise ValueError(f\"{bkv=} must be a multiple of {bkv_compute=}.\")\n  if bkv_compute % NUM_LANES:\n    raise ValueError(f\"{bkv_compute=} must be a multiple of {NUM_LANES}.\")\n\n  kv_seq_len = k.shape[kv_seq_len_dimension]\n\n  q_heads_per_kv_head = num_q_heads // num_kv_heads\n\n  if segment_ids is not None:\n    if segment_ids.q.shape != (q_seq_len,):\n      raise ValueError(\n          \"Invalid shape for q segment_ids: \"\n          f\"{segment_ids.q.shape}. Expected: {(q_seq_len,)}\"\n      )\n    if segment_ids.kv.shape != (kv_seq_len,):\n      raise ValueError(\n          \"Invalid shape for kv segment_ids: \"\n          f\"{segment_ids.kv.shape}. Expected: {(kv_seq_len,)}\"\n      )\n\n  q_layout = block_sizes.q_layout\n  def q_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None):\n    del j, data_next_ref, mask_next_ref, block_mask_ref\n    return from_head_minor((h, i, 0), q_layout)\n  def out_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None):\n    del j, data_next_ref, mask_next_ref, block_mask_ref\n    return h, i, 0\n\n  k_layout = block_sizes.k_layout","sourceCodeStart":961,"sourceCodeEnd":997,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L961-L997","documentation":"When segment_ids are provided, the query segment ids must be a 1-D array of length q_seq_len (per-sequence, no batch or head axes). Any other shape is rejected before launching the kernel.","triggerScenarios":"Passing segment_ids.q with shape [batch, q_seq_len] or [num_heads, q_seq_len] instead of [q_seq_len]; forgetting to index out the batch when using vmap with in_axes configured wrongly.","commonSituations":"Porting from GPU pipeshard/flash attention where segment ids carry a batch dim; vmap'ing attention but leaving segment ids un-vmapped.","solutions":["Pass a flat 1-D array: segment_ids.q = q_ids[q_seq_len]","If batched, vmap over segment ids too (in_axes=(..., 0, 0, SegmentIds(0, 0))) or index per batch manually","Verify kv segment ids shape is (kv_seq_len,) at the same time"],"exampleFix":"// before\nseg = SegmentIds(q=q_ids, kv=kv_ids)  # q_ids: [B, S]\nout = fn(q[0], k[0], v[0], segment_ids=seg)\n// after\nout = fn(q[0], k[0], v[0], segment_ids=SegmentIds(q_ids[0], kv_ids[0]))","handlingStrategy":"validation","validationCode":"assert segment_ids.q.shape == (q_seq_len,), f'{segment_ids.q.shape} != {(q_seq_len,)}'","typeGuard":"def q_ids_valid(ids, q_seq_len) -> bool:\n    return ids.ndim == 1 and ids.shape[0] == q_seq_len","tryCatchPattern":null,"preventionTips":["Keep segment ids 1-D per sequence","When vmapping, map segment ids with matching in_axes"],"tags":["jax","pallas","tpu","splash-attention","segment-ids","shape"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}