jax-ml/jax · error · NotImplementedError

Only 2-argument concatenate is supported.

Error message

Only 2-argument concatenate is supported.

What it means

The Triton Pallas lowering of lax.concatenate only implements the 2-argument case because it lowers to a single tt.join operation. Passing any other number of arrays raises NotImplementedError.

Source

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

  ty = ir.RankedTensorType(a.type)
  return tt_dialect.reshape(
      ir.RankedTensorType.get(shape, ty.element_type, ty.encoding),
      a,
      allow_reorder=False,
  )


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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce to pairwise joins: concatenate(a, concatenate(b, c)) as nested 2-arg calls
2. Reshape to [..., 1] per argument so the join-based lowering applies (also required by the sibling checks)
3. Move the concatenate outside the kernel
4. Use jnp.stack of 2 elements if semantics allow, since stack is also lowered via join

Example fix

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

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

def concat2(*arrays):
    assert len(arrays) <= 2 or all(a.shape[-1] == 1 for a in arrays), 'use pairwise'
    return jnp.concatenate(arrays, axis=-1)

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

try:
    out = jnp.concatenate(arrs, axis=-1)
except NotImplementedError:
    out = arrs[0]
    for a in arrs[1:]:
        out = jnp.concatenate([out, a], axis=-1)

When it happens

Trigger: Calling jnp.concatenate([...]) with fewer or more than 2 arrays inside a Triton Pallas kernel body; e.g. jnp.concatenate([a, b, c], axis=-1).

Common situations: Porting numpy/Triton code that concatenates lists of varying length; building index vectors from multiple pieces; kernels that worked on the TPU Pallas backend (which supports general concatenate).

Related errors

• Only concatenate along the last dimension is supported.
• Only arguments with shape [..., 1] are supported.
• Concatenation with Element indexing is not yet supported.
• cannot cast {src} to {dst_type}
• Only 2-argument stack is supported in Triton.

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