{"record":{"id":"aa51a7a30793b9eb","repo":"jax-ml/jax","slug":"only-2-argument-stack-is-supported-in-triton","errorCode":null,"errorMessage":"Only 2-argument stack is supported in Triton.","messagePattern":"Only 2-argument stack is supported in Triton\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1837,"sourceCode":"  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:\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):","sourceCodeStart":1819,"sourceCodeEnd":1855,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1819-L1855","documentation":"jnp.stack lowers in the Triton Pallas backend to the same tt.join primitive as concatenate, so only exactly 2 arguments can be stacked. Any other count raises NotImplementedError.","triggerScenarios":"jnp.stack([a, b, c]) or jnp.stack([a]) inside a Triton Pallas kernel body.","commonSituations":"Collecting per-iteration results into a list and stacking them inside the kernel; code reused from the TPU Pallas path where stack is fully supported.","solutions":["Nest 2-argument stacks: jnp.stack([a, jnp.stack([b, c])])","Use concatenate with expand_dims to [..., 1] args (equivalent lowering)","Move the stack out of the kernel to host code"],"exampleFix":"// before\nout = jnp.stack([a, b, c], axis=-1)\n\n// after\nout = jnp.stack([a, jnp.stack([b, c], axis=-1)], axis=-1)","handlingStrategy":"validation","validationCode":"def stack2(arrs, axis=-1):\n    out = arrs[-1]\n    for a in reversed(arrs[:-1]):\n        out = jnp.stack([a, out], axis=axis)\n    return out","typeGuard":"def is_pair_stack(arrs) -> bool:\n    return len(arrs) == 2","tryCatchPattern":"try:\n    out = jnp.stack(arrs, axis=-1)\nexcept NotImplementedError:\n    out = stack2(arrs, axis=-1)","preventionTips":["Only stack exactly 2 arrays in Triton Pallas kernels","Fold N-way stacks into nested 2-way helper","Prefer keeping values as separate kernel arguments"],"tags":["jax","triton","pallas","stack","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"}