jax-ml/jax · error · NotImplementedError

Only concatenate along the last dimension is supported.

Error message

Only concatenate along the last dimension is supported.

What it means

Triton's tt.join only combines two scalars along a new minor dimension, so the Pallas lowering only supports concatenate along the last dimension. Any other dimension raises NotImplementedError.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Transpose so the concatenation axis is last, concatenate, then transpose back
Move the concatenate outside the kernel (host-side jnp.concatenate)
Reorder the kernel's block layout so the joined axis is the minor dimension

Example fix

// before
out = jnp.concatenate([x, y], axis=0)  # 2D blocks

// after
out = jnp.transpose(jnp.concatenate([jnp.transpose(x), jnp.transpose(y)], axis=-1))

Defensive patterns

Strategy: validation

Validate before calling

def safe_concat(x, y, axis):
    if axis != x.ndim - 1:
        perm = tuple(range(x.ndim - 1)) if False else None
        raise ValueError('transpose to last-axis concat inside kernels')
    return jnp.concatenate([x, y], axis=axis)

Type guard

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

Try / catch

try:
    out = jnp.concatenate([x, y], axis=axis)
except NotImplementedError:
    xt, yt = jnp.transpose(x), jnp.transpose(y)
    out = jnp.transpose(jnp.concatenate([xt, yt], axis=-1))

Prevention

Standardize on axis=-1 for in-kernel concatenation
Encapsulate concat in kernel helpers that assert the axis
Keep axis logic in kernel wrappers, not the kernel body

When it happens

Trigger: jnp.concatenate([x, y], axis=k) inside a Triton Pallas kernel where k is not x.ndim-1, e.g. concatenating along axis 0 of a 2D block.

Common situations: Kernels ported from XLA/TPU Pallas or numpy that concatenate along the leading (batch) axis; assembling coordinate vectors along a non-minor axis.

Related errors

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