jax-ml/jax · error · NotImplementedError

Only stack along the last dimension is supported in Triton.

Error message

Only stack along the last dimension is supported in Triton.

What it means

Stack is lowered to tt.join, which appends a new minor dimension; therefore the Triton Pallas backend only supports stacking along a brand-new last dimension (axis == x.ndim). Stacking along any existing axis is unimplemented.

Source

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

        "Only concatenate along the last dimension is supported."
    )
  if x_aval.shape[-1] != 1 or y_aval.shape[-1] != 1:
    raise NotImplementedError(
        "Only arguments with shape [..., 1] are supported."
    )
  lhs = _reshape(x, x_aval.shape[:-1])
  rhs = _reshape(y, y_aval.shape[:-1])
  ret_type = get_join_type(ir.RankedTensorType(rhs.type))
  return tt_dialect.join(ret_type, lhs, rhs)

@register_lowering(jax._src.lax.lax.stack_p)
def _stack_lowering_rule(ctx: LoweringRuleContext, *args, axis):
  if len(args) != 2:
    raise NotImplementedError("Only 2-argument stack is supported in Triton.")
  [x_aval, y_aval] = ctx.avals_in
  x, y = args
  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass axis=-1 (equivalently axis=x.ndim) when stacking inside the kernel
2. Use expand_dims + concatenate instead for other axes
3. Transpose after an axis=-1 stack to move the new dimension where needed

Example fix

// before
out = jnp.stack([x, y])  # default axis=0

# after
out = jnp.stack([x, y], axis=-1)

def safe_stack(x, y, axis=None):
    axis = x.ndim if axis in (None, -1) else axis
    assert axis == x.ndim, 'in-kernel stack must use a new last axis'
    return jnp.stack([x, y], axis=axis)

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

try:
    out = jnp.stack([x, y], axis=0)
except NotImplementedError:
    out = jnp.stack([x, y], axis=-1)

When it happens

Trigger: jnp.stack([x, y], axis=k) inside a Triton Pallas kernel where k != x.ndim, e.g. stacking 1D arrays along axis 0.

Common situations: Default jnp.stack behavior (axis=0) used inside kernels — since default is 0 and ndim >= 1, 1D stacks with default axis hit this immediately; ported numpy logic.

Related errors

• Only concatenate along the last dimension is supported.
• Only 2-argument stack is supported in Triton.
• Only unstack along the last dimension is supported in Triton
• Stack with Element indexing is not yet supported.
• cannot cast {src} to {dst_type}

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