jax-ml/jax · error · NotImplementedError

Only unstack along the last dimension is supported in Triton

Error message

Only unstack along the last dimension is supported in Triton.

What it means

tt.split always splits the minor (last) dimension, so the Triton Pallas lowering of unstack requires axis == x.ndim - 1. Unstacking along any other axis is not implemented.

Source

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

  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):
    shape = ir.RankedTensorType(x.type).shape
    x = _reshape(x, shape[:axis] + [2, shape[axis] // 2] + shape[axis + 1 :])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose first so the size-2 axis is last, unstack, then proceed
2. Store the pair along the last dimension when constructing the tensor
3. Use explicit indexing v[..., 0], v[..., 1] instead of unstack

Example fix

// before
a, b = jax.lax.unstack(v, axis=0)  # v shape [2, B]

// after
a, b = jax.lax.unstack(jnp.transpose(v), axis=-1)

def safe_unstack(v, axis=-1):
    if axis != v.ndim - 1:
        v = jnp.moveaxis(v, axis, -1)
    return jax.lax.unstack(v, axis=-1)

def last_axis_split_ok(axis, ndim) -> bool:
    return axis == ndim - 1

try:
    parts = jax.lax.unstack(v, axis=axis)
except NotImplementedError:
    parts = jax.lax.unstack(jnp.moveaxis(v, axis, -1), axis=-1)

When it happens

Trigger: 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.

Common situations: Unstacking paired values stored along the leading axis; code written against TPU Pallas where any axis works.

Related errors

• Only concatenate along the last dimension is supported.
• Only stack along the last dimension is supported in Triton.
• Only unstack of size 2 is supported in Triton.
• cannot cast {src} to {dst_type}
• Only 2-argument concatenate is supported.

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