jax-ml/jax · error · NotImplementedError

Only unstack of size 2 is supported in Triton.

Error message

Only unstack of size 2 is supported in Triton.

What it means

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.

Source

Thrown at jax/_src/pallas/triton/lowering.py:1858

  if axis != x_aval.ndim:
    raise NotImplementedError("Only stack along the last dimension is supported in Triton.")

  x = _ensure_ir_value(x, x_aval)
  y = _ensure_ir_value(y, y_aval)

  ty = ir.RankedTensorType(x.type)
  shape = list(ty.shape)
  shape.append(2)
  ret_type = ir.RankedTensorType.get(shape, ty.element_type, ty.encoding)

  return tt_dialect.join(ret_type, x, y)


@register_lowering(jax._src.lax.lax.unstack_p)
def _unstack_lowering_rule(ctx: LoweringRuleContext, x, *, axis):
  [x_aval] = ctx.avals_in
  if x_aval.shape[axis] != 2:
    raise NotImplementedError("Only unstack of size 2 is supported in Triton.")
  if axis != x_aval.ndim - 1:
    raise NotImplementedError("Only unstack along the last dimension is supported in Triton.")

  x = _ensure_ir_value(x, x_aval)
  return tuple(tt_dialect.split(x))


@register_lowering(lax.split_p)
def _split_lowering_rule(ctx: LoweringRuleContext, x, *, sizes, axis):
  pass
  # TODO(cjfj): Add support for larger powers of 2.
  num_parts = len(sizes)
  if num_parts != pallas_utils.next_power_of_2(num_parts):
    raise NotImplementedError("Only power-of-2 num parts supported.")
  if any(size != sizes[0] for size in sizes):
    raise NotImplementedError("Only equal-sized splits are supported.")

  def split_into_2(x):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use slicing (v[0], v[1], ... or lax.index_in_dim) instead of unstack for sizes != 2
2. Restructure the kernel to keep components as separate arguments rather than one stacked tensor
3. Nest unstacks if the size is a power of 2 via split

Example fix

// before
a, b, c = jax.lax.unstack(v)  # v.shape[axis] == 3

// after
a, b, c = v[0], v[1], v[2]

assert v.shape[axis] == 2, 'in-kernel unstack requires axis size 2'

def unstackable(v, axis=-1) -> bool:
    return v.shape[axis] == 2 and axis == v.ndim - 1

try:
    a, b = jax.lax.unstack(v, axis=-1)
except NotImplementedError:
    a, b = v[..., 0], v[..., 1]

When it happens

Trigger: x, y = jax.lax.unstack(v) where v.shape[axis] != 2 (e.g. 3 or more stacked elements) inside a Triton Pallas kernel.

Common situations: Symmetric pairing code works (size 2) but generalized to N elements; iterating over stacked results with unpacking syntax inside kernels.

Related errors

• Only unstack along the last dimension is supported in Triton
• cannot cast {src} to {dst_type}
• Only 2-argument concatenate is supported.
• Only concatenate along the last dimension is supported.
• Only arguments with shape [..., 1] are supported.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/a529d3df53f21586. Report an issue: GitHub.