{"record":{"id":"dd81ae95496be204","repo":"jax-ml/jax","slug":"only-2-argument-concatenate-is-supported","errorCode":null,"errorMessage":"Only 2-argument concatenate is supported.","messagePattern":"Only 2-argument concatenate is supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1818,"sourceCode":"\n  ty = ir.RankedTensorType(a.type)\n  return tt_dialect.reshape(\n      ir.RankedTensorType.get(shape, ty.element_type, ty.encoding),\n      a,\n      allow_reorder=False,\n  )\n\n\ndef get_join_type(old_type: ir.RankedTensorType):\n  shape = old_type.shape\n  shape.append(2)\n  return ir.RankedTensorType.get(shape, old_type.element_type, old_type.encoding)\n\n\n@register_lowering(lax.concatenate_p)\ndef _concatenate_lowering_rule(ctx: LoweringRuleContext, *args, dimension):\n  if len(args) != 2:\n    raise NotImplementedError(\"Only 2-argument concatenate is supported.\")\n  x_aval, y_aval = ctx.avals_in\n  x, y = args\n  if dimension != x_aval.ndim-1:\n    raise NotImplementedError(\n        \"Only concatenate along the last dimension is supported.\"\n    )\n  if x_aval.shape[-1] != 1 or y_aval.shape[-1] != 1:\n    raise NotImplementedError(\n        \"Only arguments with shape [..., 1] are supported.\"\n    )\n  lhs = _reshape(x, x_aval.shape[:-1])\n  rhs = _reshape(y, y_aval.shape[:-1])\n  ret_type = get_join_type(ir.RankedTensorType(rhs.type))\n  return tt_dialect.join(ret_type, lhs, rhs)\n\n@register_lowering(jax._src.lax.lax.stack_p)\ndef _stack_lowering_rule(ctx: LoweringRuleContext, *args, axis):\n  if len(args) != 2:","sourceCodeStart":1800,"sourceCodeEnd":1836,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1800-L1836","documentation":"The Triton Pallas lowering of lax.concatenate only implements the 2-argument case because it lowers to a single tt.join operation. Passing any other number of arrays raises NotImplementedError.","triggerScenarios":"Calling jnp.concatenate([...]) with fewer or more than 2 arrays inside a Triton Pallas kernel body; e.g. jnp.concatenate([a, b, c], axis=-1).","commonSituations":"Porting numpy/Triton code that concatenates lists of varying length; building index vectors from multiple pieces; kernels that worked on the TPU Pallas backend (which supports general concatenate).","solutions":["Reduce to pairwise joins: concatenate(a, concatenate(b, c)) as nested 2-arg calls","Reshape to [..., 1] per argument so the join-based lowering applies (also required by the sibling checks)","Move the concatenate outside the kernel","Use jnp.stack of 2 elements if semantics allow, since stack is also lowered via join"],"exampleFix":"// before\nout = jnp.concatenate([a, b, c], axis=-1)\n\n// after\nout = jnp.concatenate([a, jnp.concatenate([b, c], axis=-1)], axis=-1)","handlingStrategy":"validation","validationCode":"def concat2(*arrays):\n    assert len(arrays) <= 2 or all(a.shape[-1] == 1 for a in arrays), 'use pairwise'\n    return jnp.concatenate(arrays, axis=-1)","typeGuard":"def is_pairwise_concat(arrs) -> bool:\n    return len(arrs) == 2","tryCatchPattern":"try:\n    out = jnp.concatenate(arrs, axis=-1)\nexcept NotImplementedError:\n    out = arrs[0]\n    for a in arrs[1:]:\n        out = jnp.concatenate([out, a], axis=-1)","preventionTips":["Always concatenate exactly 2 arrays inside Triton Pallas kernels","Wrap concatenation in a helper that folds to pairwise 2-arg calls","Move general concatenation to host code"],"tags":["jax","triton","pallas","concatenate","not-implemented"],"backgroundTag":"unsupported-operation-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}