{"record":{"id":"79055c2130387551","repo":"jax-ml/jax","slug":"only-concatenate-along-the-last-dimension-is-suppo","errorCode":null,"errorMessage":"Only concatenate along the last dimension is supported.","messagePattern":"Only concatenate along the last dimension is supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1822,"sourceCode":"      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:\n    raise NotImplementedError(\"Only 2-argument stack is supported in Triton.\")\n  [x_aval, y_aval] = ctx.avals_in\n  x, y = args\n  if axis != x_aval.ndim:","sourceCodeStart":1804,"sourceCodeEnd":1840,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1804-L1840","documentation":"Triton's tt.join only combines two scalars along a new minor dimension, so the Pallas lowering only supports concatenate along the last dimension. Any other dimension raises NotImplementedError.","triggerScenarios":"jnp.concatenate([x, y], axis=k) inside a Triton Pallas kernel where k is not x.ndim-1, e.g. concatenating along axis 0 of a 2D block.","commonSituations":"Kernels ported from XLA/TPU Pallas or numpy that concatenate along the leading (batch) axis; assembling coordinate vectors along a non-minor axis.","solutions":["Transpose so the concatenation axis is last, concatenate, then transpose back","Move the concatenate outside the kernel (host-side jnp.concatenate)","Reorder the kernel's block layout so the joined axis is the minor dimension"],"exampleFix":"// before\nout = jnp.concatenate([x, y], axis=0)  # 2D blocks\n\n// after\nout = jnp.transpose(jnp.concatenate([jnp.transpose(x), jnp.transpose(y)], axis=-1))","handlingStrategy":"validation","validationCode":"def safe_concat(x, y, axis):\n    if axis != x.ndim - 1:\n        perm = tuple(range(x.ndim - 1)) if False else None\n        raise ValueError('transpose to last-axis concat inside kernels')\n    return jnp.concatenate([x, y], axis=axis)","typeGuard":"def last_axis(axis, ndim) -> bool:\n    return axis == ndim - 1","tryCatchPattern":"try:\n    out = jnp.concatenate([x, y], axis=axis)\nexcept NotImplementedError:\n    xt, yt = jnp.transpose(x), jnp.transpose(y)\n    out = jnp.transpose(jnp.concatenate([xt, yt], axis=-1))","preventionTips":["Standardize on axis=-1 for in-kernel concatenation","Encapsulate concat in kernel helpers that assert the axis","Keep axis logic in kernel wrappers, not the kernel body"],"tags":["jax","triton","pallas","concatenate","axis","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"}