jax-ml/jax · error · NotImplementedError

Only unstack along the last dimension is supported in…

Error message

Only unstack along the last dimension is supported in Triton.

What it means

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.

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 to
• loading from a block pointer is not supported

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

Appendix: 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)