{"record":{"id":"2d629aca15f178f9","repo":"jax-ml/jax","slug":"concatenation-with-element-indexing-is-not-yet-sup","errorCode":null,"errorMessage":"Concatenation with Element indexing is not yet supported.","messagePattern":"Concatenation with Element indexing is not yet supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":1742,"sourceCode":"    if idx_aval.start != 0:\n      raise NotImplementedError('get not supported yet')\n    if idx_aval.size != size:\n      raise NotImplementedError('get not supported yet')\n    bidx = next(block_idx_iter)\n    block_indexer.append(_slice(bidx, bd))\n  assert next(block_idx_iter, None) is None\n  return ref.get(idx=tuple(block_indexer))\n\n\n@register_eval_rule(lax.concatenate_p)\ndef _concatenate_eval_rule(ctx: KernelEvalContext, *args, dimension):\n  # We now handle the case where each of the concatenated array dimensions\n  # divides the block size.\n  block_spec = ctx.out_block_specs[0]\n  block_shape = block_spec.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        'Concatenation with Element indexing is not yet supported.'\n    )\n  block_dim = block_shape[dimension]\n  if block_dim is None:\n    block_dim = 1\n\n  if block_dim == sum(aval.shape[dimension] for aval in ctx.avals_in):\n    # Handle special case if the block contains all of the concatenated\n    # array.\n    return jax.lax.concatenate(args, dimension=dimension)\n\n  num_blocks = []\n  for aval in ctx.avals_in:\n    assert isinstance(aval, core.ShapedArray)\n    if aval.shape[dimension] % block_dim != 0:\n      raise ValueError(\n          f'Shape along concat dimension {dimension} must be divisible by the'\n          f' block shape {block_shape[dimension]} for all children. Got shape'","sourceCodeStart":1724,"sourceCodeEnd":1760,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L1724-L1760","documentation":"Raised by _concatenate_eval_rule when concatenation is evaluated on a Ref whose output BlockSpec contains pallas_core.Element block entries. The fuser's concat strategy relies on counting whole blocks along the concat dimension, which is meaningless for Element indexing.","triggerScenarios":"Using jnp.concatenate / lax.concatenate on values that flow into a pallas output ref blocked with pallas_core.Element along any dimension.","commonSituations":"Building outputs by concatenating tiles inside kernels with fine-grained Element block specs; upgrading JAX where Element blocks became available but concat support lagged.","solutions":["Remove Element entries from the output BlockSpec; use int block sizes","Compute the concatenation outside the kernel and write blocks individually","Use jnp.stack-free manual assignment of slices into the output block"],"exampleFix":"# before\nout_ref[...] = jnp.concatenate([a, b], axis=0)  # out uses Element blocks\n\n# after\nout_ref[...] = jnp.concatenate([a, b], axis=0)  # BlockSpec uses int blocks\n# BlockSpec(block_shape=(BS, ...)) instead of (Element(...), ...)","handlingStrategy":"validation","validationCode":"import jax._src.pallas.pallas_core as pc\nassert not any(isinstance(b, pc.Element) for b in block_spec.block_shape), 'concat output cannot use Element blocks'","typeGuard":"def concat_safe_block_spec(bs):\n    import jax._src.pallas.pallas_core as pc\n    return not any(isinstance(b, pc.Element) for b in bs.block_shape)","tryCatchPattern":null,"preventionTips":["Avoid Element blocks on outputs receiving concatenated values","Write blocks individually instead of concatenating into Element-blocked refs"],"tags":["jax","pallas","concatenate","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"}