{"record":{"id":"9fec7a095a5962d1","repo":"jax-ml/jax","slug":"stack-with-element-indexing-is-not-yet-supported","errorCode":null,"errorMessage":"Stack with Element indexing is not yet supported.","messagePattern":"Stack with Element indexing is not yet supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":1866,"sourceCode":"      return util.tuple_update(idx, dimension, block_idx)\n\n    return block_transform.replace(\n        block_index_transform=new_block_index_transform\n    )\n  return [make_block_transform(i) for i in range(len(ctx.avals_in))]\n\n\n@register_pull_block_spec_rule(lax.stack_p)\ndef _stack_pull_rule(\n    ctx: PullRuleContext,\n    block_transform: BlockIndexTransform,\n    *,\n    axis: int,\n):\n  block_shape = block_transform.block_shape\n  is_element_block = [isinstance(bd, pallas_core.Element) for bd in block_shape]\n  if any(is_element_block):\n    raise NotImplementedError(\n        'Stack with Element indexing is not yet supported.'\n    )\n  block_dim = block_shape[axis]\n  if block_dim is None or isinstance(block_dim, pallas_core.Squeezed):\n    block_dim = 1\n\n  n = len(ctx.avals_in)\n  if block_dim != n:\n    raise NotImplementedError(\n        \"Stacking only supported when the block size along the stack axis \"\n        f\"equals the number of inputs. Got block_dim={block_dim}, expected {n}.\"\n    )\n\n  new_block_shape = list(block_transform.block_shape)\n  new_block_shape.pop(axis)\n\n  def make_block_transform(child_index: int):\n    def new_block_index_transform(*idxs):","sourceCodeStart":1848,"sourceCodeEnd":1884,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L1848-L1884","documentation":"Raised by _stack_pull_rule when propagating block specs through jnp.stack where the output block transform uses pallas_core.Element entries. Stack's block decomposition (block size along the new axis equals number of inputs) cannot be expressed with Element indexing.","triggerScenarios":"Using jnp.stack inside a pallas kernel where the stacked result feeds a ref whose BlockSpec uses pallas_core.Element along any dimension.","commonSituations":"Assembling per-block stacked outputs with fine-grained Element layouts; combining Mosaic Element indexing with stack-based output assembly.","solutions":["Replace Element block entries with int/Squeezed/None descriptors in the output BlockSpec","Use concatenate or per-input writes instead of stack","Perform the stack outside the kernel and pass the result as a normally-blocked input"],"exampleFix":"# before\nBlockSpec(block_shape=(pallas_core.Element(1), 64), ...)\nout[...] = jnp.stack([a, b], axis=0)\n\n# after\nBlockSpec(block_shape=(2, 64), ...)\nout[...] = jnp.stack([a, b], axis=0)","handlingStrategy":"validation","validationCode":"import jax._src.pallas.pallas_core as pc\nassert not any(isinstance(b, pc.Element) for b in block_shape), 'stack output cannot use Element blocks'","typeGuard":"def stack_safe_block_spec(block_shape):\n    import jax._src.pallas.pallas_core as pc\n    return not any(isinstance(b, pc.Element) for b in block_shape)","tryCatchPattern":null,"preventionTips":["Use int block sizes (equal to number of stacked inputs along the new axis) instead of Element","Prefer per-input writes or outside-kernel stacking"],"tags":["jax","pallas","stack","element-indexing","not-implemented"],"backgroundTag":"jax-pallas-unsupported-blockspec-op","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}