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

The Triton concatenate lowering relies on tt.join, which operates on the minor (last) dimension, so the concatenation axis must be ndim-1. Concatenating along any non-minor dimension is unimplemented and fails in the lowering rule.

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.

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