{"record":{"id":"5a8f9b07e62532d4","repo":"keras-team/keras","slug":"sharding-along-sequence-dimension-not-allowed-in-t-5a8f9b","errorCode":null,"errorMessage":"Sharding along sequence dimension not allowed in TPU kernel attention","messagePattern":"Sharding along sequence dimension not allowed in TPU kernel attention","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"keras/src/backend/jax/ops/nn.py","lineNumber":1613,"sourceCode":"        decoder_segment_ids: Optional. Segment IDs for the decoder, used for\n            sharding or masking.\n        custom_mask: Optional[jax.Array]. A custom attention mask to apply. If\n            None, a causal mask is used.\n        attn_logits_soft_cap: Optional[float]. If provided, applies a soft cap\n            to the attention logits.\n        head_shards: int, default=1. Number of shards for the attention heads.\n        q_seq_shards: int, default=1. Number of shards for the query sequence\n            dimension.\n    Returns:\n        jax.Array: The result of applying the Splash multi-head attention\n            kernel to the inputs.\n    Raises:\n        AssertionError: If sharding along the sequence dimension is attempted\n            with decoder_segment_ids.\n    \"\"\"\n    if decoder_segment_ids is not None:\n        if query.shape[2] != decoder_segment_ids.q.shape[1]:\n            raise ValueError(\n                \"Sharding along sequence dimension not allowed\"\n                \" in TPU kernel attention\"\n            )\n\n    if custom_mask is not None:\n        mask = splash_attention_mask.NumpyMask(array=custom_mask)\n    else:\n        mask = splash_attention_mask.CausalMask(\n            shape=(query.shape[2], query.shape[2])\n        )\n\n    # Create multi-head mask\n    multi_head_mask = splash_attention_mask.MultiHeadMask(\n        masks=(mask,) * query.shape[1]\n    )\n    splash_kernel = splash_attention_kernel.make_splash_mha(\n        mask=multi_head_mask,\n        head_shards=head_shards,","sourceCodeStart":1595,"sourceCodeEnd":1631,"githubUrl":"https://github.com/keras-team/keras/blob/7a34a03db60bf60042242d6a556fc3be119046a5/keras/src/backend/jax/ops/nn.py#L1595-L1631","documentation":"Error \"Sharding along sequence dimension not allowed in TPU kernel attention\" thrown in keras-team/keras.","triggerScenarios":"Thrown at keras/src/backend/jax/ops/nn.py:1613 when the library encounters an invalid state.","commonSituations":"See trigger scenarios.","solutions":[],"exampleFix":null,"handlingStrategy":null,"validationCode":null,"typeGuard":null,"tryCatchPattern":null,"preventionTips":[],"tags":[],"backgroundTag":null,"analyzedSha":"7a34a03db60bf60042242d6a556fc3be119046a5","analyzedAt":"2026-08-25T21:25:25.994Z","schemaVersion":2},"datasetVersion":"2026-08-26T02:17:13.382Z"}