{"record":{"id":"a529d3df53f21586","repo":"jax-ml/jax","slug":"only-unstack-of-size-2-is-supported-in-triton","errorCode":null,"errorMessage":"Only unstack of size 2 is supported in Triton.","messagePattern":"Only unstack of size 2 is supported in Triton\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1858,"sourceCode":"  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:\n    raise NotImplementedError(\"Only unstack along the last dimension is supported in Triton.\")\n\n  x = _ensure_ir_value(x, x_aval)\n  return tuple(tt_dialect.split(x))\n\n\n@register_lowering(lax.split_p)\ndef _split_lowering_rule(ctx: LoweringRuleContext, x, *, sizes, axis):\n  pass\n  # TODO(cjfj): Add support for larger powers of 2.\n  num_parts = len(sizes)\n  if num_parts != pallas_utils.next_power_of_2(num_parts):\n    raise NotImplementedError(\"Only power-of-2 num parts supported.\")\n  if any(size != sizes[0] for size in sizes):\n    raise NotImplementedError(\"Only equal-sized splits are supported.\")\n\n  def split_into_2(x):","sourceCodeStart":1840,"sourceCodeEnd":1876,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1840-L1876","documentation":"Unstack lowers to tt.split, which always splits a dimension of exactly 2 into two results. The Triton Pallas backend therefore only supports unstacking an axis whose size is 2.","triggerScenarios":"x, y = jax.lax.unstack(v) where v.shape[axis] != 2 (e.g. 3 or more stacked elements) inside a Triton Pallas kernel.","commonSituations":"Symmetric pairing code works (size 2) but generalized to N elements; iterating over stacked results with unpacking syntax inside kernels.","solutions":["Use slicing (v[0], v[1], ... or lax.index_in_dim) instead of unstack for sizes != 2","Restructure the kernel to keep components as separate arguments rather than one stacked tensor","Nest unstacks if the size is a power of 2 via split"],"exampleFix":"// before\na, b, c = jax.lax.unstack(v)  # v.shape[axis] == 3\n\n// after\na, b, c = v[0], v[1], v[2]","handlingStrategy":"validation","validationCode":"assert v.shape[axis] == 2, 'in-kernel unstack requires axis size 2'","typeGuard":"def unstackable(v, axis=-1) -> bool:\n    return v.shape[axis] == 2 and axis == v.ndim - 1","tryCatchPattern":"try:\n    a, b = jax.lax.unstack(v, axis=-1)\nexcept NotImplementedError:\n    a, b = v[..., 0], v[..., 1]","preventionTips":["Use indexing instead of unstack for sizes other than 2","Keep paired values on a trailing size-2 axis","Pass components as separate kernel args instead"],"tags":["jax","triton","pallas","unstack","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"}