jax-ml/jax · error · NotImplementedError

Only arguments with shape [..., 1] are supported.

Error message

Only arguments with shape [..., 1] are supported.

What it means

The lowering of concatenate to tt.join requires each argument to have shape [..., 1] — i.e. the last dimension must be exactly 1 so it can be reshaped away and re-joined as a new minor dimension of size 2. Any other trailing shape fails.

Source

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

def get_join_type(old_type: ir.RankedTensorType):
  shape = old_type.shape
  shape.append(2)
  return ir.RankedTensorType.get(shape, old_type.element_type, old_type.encoding)


@register_lowering(lax.concatenate_p)
def _concatenate_lowering_rule(ctx: LoweringRuleContext, *args, dimension):
  if len(args) != 2:
    raise NotImplementedError("Only 2-argument concatenate is supported.")
  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restructure to pair singleton trailing dims: expand_dims each operand to [..., 1] before concatenating
2. Do general concatenation outside the kernel on the host/XLA side
3. Build the combined block with explicit indexing/store into a preallocated reference instead of concatenate

Example fix

// before
out = jnp.concatenate([x, y], axis=-1)  # x, y have shape [B, k]

// after
out = jnp.concatenate([x[..., None, :], y[..., None, :]], axis=-2)  # or concatenate outside kernel

assert x.shape[-1] == 1 and y.shape[-1] == 1, 'operands must be [..., 1] for in-kernel concat'

def singleton_trailing(a) -> bool:
    return a.shape[-1] == 1

try:
    out = jnp.concatenate([x, y], axis=-1)
except NotImplementedError:
    out = jnp.stack([x, y], axis=-1)  # if pairing blocks

When it happens

Trigger: jnp.concatenate([x, y], axis=-1) inside a Triton Pallas kernel where x.shape[-1] or y.shape[-1] is not 1, e.g. joining two [B, N] blocks with N > 1.

Common situations: Trying to append a column of values to a block (shape [B, k] + [B, k]) rather than pairing scalars; generic concatenation assumed from numpy semantics.

Related errors

• Only 2-argument concatenate is supported.
• Only concatenate along the last dimension is supported.
• Only equal-sized splits are supported.
• Concatenation 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/b2973dfa876124f7. Report an issue: GitHub.