jax-ml/jax · error · NotImplementedError

Only 2-argument stack is supported in Triton.

Error message

Only 2-argument stack is supported in Triton.

What it means

jnp.stack lowers in the Triton Pallas backend to the same tt.join primitive as concatenate, so only exactly 2 arguments can be stacked. Any other count raises NotImplementedError.

Source

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

  x_aval, y_aval = ctx.avals_in
  x, y = args
  if dimension != x_aval.ndim-1:
    raise NotImplementedError(
        "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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Nest 2-argument stacks: jnp.stack([a, jnp.stack([b, c])])
2. Use concatenate with expand_dims to [..., 1] args (equivalent lowering)
3. Move the stack out of the kernel to host code

Example fix

// before
out = jnp.stack([a, b, c], axis=-1)

// after
out = jnp.stack([a, jnp.stack([b, c], axis=-1)], axis=-1)

def stack2(arrs, axis=-1):
    out = arrs[-1]
    for a in reversed(arrs[:-1]):
        out = jnp.stack([a, out], axis=axis)
    return out

def is_pair_stack(arrs) -> bool:
    return len(arrs) == 2

try:
    out = jnp.stack(arrs, axis=-1)
except NotImplementedError:
    out = stack2(arrs, axis=-1)

When it happens

Trigger: jnp.stack([a, b, c]) or jnp.stack([a]) inside a Triton Pallas kernel body.

Common situations: Collecting per-iteration results into a list and stacking them inside the kernel; code reused from the TPU Pallas path where stack is fully supported.

Related errors

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

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