{"record":{"id":"58267fe522f4f4dc","repo":"jax-ml/jax","slug":"only-unstack-along-the-last-dimension-is-supported","errorCode":null,"errorMessage":"Only unstack along the last dimension is supported in Triton.","messagePattern":"Only unstack along the last dimension is supported in Triton\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1860,"sourceCode":"\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):\n    shape = ir.RankedTensorType(x.type).shape\n    x = _reshape(x, shape[:axis] + [2, shape[axis] // 2] + shape[axis + 1 :])","sourceCodeStart":1842,"sourceCodeEnd":1878,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1842-L1878","documentation":"The Triton lowering of unstack is only implemented for the trailing (last) dimension, matching Triton's register layout. It fires when the unstack axis is not the final axis of the array.","triggerScenarios":"jax.lax.unstack(v, axis=k) with k != v.ndim - 1 inside a Triton Pallas kernel, e.g. unstacking a [2, B] tensor along axis 0.","commonSituations":"Unstacking paired values stored along the leading axis; code written against TPU Pallas where any axis works.","solutions":["Transpose first so the size-2 axis is last, unstack, then proceed","Store the pair along the last dimension when constructing the tensor","Use explicit indexing v[..., 0], v[..., 1] instead of unstack"],"exampleFix":"// before\na, b = jax.lax.unstack(v, axis=0)  # v shape [2, B]\n\n// after\na, b = jax.lax.unstack(jnp.transpose(v), axis=-1)","handlingStrategy":"validation","validationCode":"def safe_unstack(v, axis=-1):\n    if axis != v.ndim - 1:\n        v = jnp.moveaxis(v, axis, -1)\n    return jax.lax.unstack(v, axis=-1)","typeGuard":"def last_axis_split_ok(axis, ndim) -> bool:\n    return axis == ndim - 1","tryCatchPattern":"try:\n    parts = jax.lax.unstack(v, axis=axis)\nexcept NotImplementedError:\n    parts = jax.lax.unstack(jnp.moveaxis(v, axis, -1), axis=-1)","preventionTips":["Always unstack along the last axis in Triton Pallas kernels","moveaxis before unstack when data layout differs","Prefer explicit indexing for one-off extraction"],"tags":["jax","triton","pallas","unstack","axis","not-implemented"],"backgroundTag":"unsupported-operation-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","contentChangedAt":null,"schemaVersion":2},"datasetVersion":"2026-09-23T08:17:48.524Z"}