{"record":{"id":"ea555a2a150f5b20","repo":"jax-ml/jax","slug":"only-stack-along-the-last-dimension-is-supported-i","errorCode":null,"errorMessage":"Only stack along the last dimension is supported in Triton.","messagePattern":"Only stack along the last dimension is supported in Triton\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1841,"sourceCode":"        \"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:\n    raise NotImplementedError(\"Only stack along the last dimension is supported in Triton.\")\n\n  x = _ensure_ir_value(x, x_aval)\n  y = _ensure_ir_value(y, y_aval)\n\n  ty = ir.RankedTensorType(x.type)\n  shape = list(ty.shape)\n  shape.append(2)\n  ret_type = ir.RankedTensorType.get(shape, ty.element_type, ty.encoding)\n\n  return tt_dialect.join(ret_type, x, y)\n\n\n@register_lowering(jax._src.lax.lax.unstack_p)\ndef _unstack_lowering_rule(ctx: LoweringRuleContext, x, *, axis):\n  [x_aval] = ctx.avals_in\n  if x_aval.shape[axis] != 2:\n    raise NotImplementedError(\"Only unstack of size 2 is supported in Triton.\")\n  if axis != x_aval.ndim - 1:","sourceCodeStart":1823,"sourceCodeEnd":1859,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1823-L1859","documentation":"Stack is lowered to tt.join, which appends a new minor dimension; therefore the Triton Pallas backend only supports stacking along a brand-new last dimension (axis == x.ndim). Stacking along any existing axis is unimplemented.","triggerScenarios":"jnp.stack([x, y], axis=k) inside a Triton Pallas kernel where k != x.ndim, e.g. stacking 1D arrays along axis 0.","commonSituations":"Default jnp.stack behavior (axis=0) used inside kernels — since default is 0 and ndim >= 1, 1D stacks with default axis hit this immediately; ported numpy logic.","solutions":["Pass axis=-1 (equivalently axis=x.ndim) when stacking inside the kernel","Use expand_dims + concatenate instead for other axes","Transpose after an axis=-1 stack to move the new dimension where needed"],"exampleFix":"// before\nout = jnp.stack([x, y])  # default axis=0\n\n# after\nout = jnp.stack([x, y], axis=-1)","handlingStrategy":"validation","validationCode":"def safe_stack(x, y, axis=None):\n    axis = x.ndim if axis in (None, -1) else axis\n    assert axis == x.ndim, 'in-kernel stack must use a new last axis'\n    return jnp.stack([x, y], axis=axis)","typeGuard":"def new_last_axis(axis, ndim) -> bool:\n    return axis == ndim or axis == -1","tryCatchPattern":"try:\n    out = jnp.stack([x, y], axis=0)\nexcept NotImplementedError:\n    out = jnp.stack([x, y], axis=-1)","preventionTips":["Always pass axis=-1 to stack inside kernels (default 0 fails)","Transpose afterwards to relocate the new dimension","Lint kernel code for jnp.stack without axis=-1"],"tags":["jax","triton","pallas","stack","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"}