jax-ml/jax · error · NotImplementedError

Concatenation with Element indexing is not yet supported.

Error message

Concatenation with Element indexing is not yet supported.

What it means

Raised by _concatenate_eval_rule when concatenation is evaluated on a Ref whose output BlockSpec contains pallas_core.Element block entries. The fuser's concat strategy relies on counting whole blocks along the concat dimension, which is meaningless for Element indexing.

Source

Thrown at jax/_src/pallas/fuser/block_spec.py:1742

    if idx_aval.start != 0:
      raise NotImplementedError('get not supported yet')
    if idx_aval.size != size:
      raise NotImplementedError('get not supported yet')
    bidx = next(block_idx_iter)
    block_indexer.append(_slice(bidx, bd))
  assert next(block_idx_iter, None) is None
  return ref.get(idx=tuple(block_indexer))


@register_eval_rule(lax.concatenate_p)
def _concatenate_eval_rule(ctx: KernelEvalContext, *args, dimension):
  # We now handle the case where each of the concatenated array dimensions
  # divides the block size.
  block_spec = ctx.out_block_specs[0]
  block_shape = block_spec.block_shape
  is_element_block = [isinstance(bd, pallas_core.Element) for bd in block_shape]
  if any(is_element_block):
    raise NotImplementedError(
        'Concatenation with Element indexing is not yet supported.'
    )
  block_dim = block_shape[dimension]
  if block_dim is None:
    block_dim = 1

  if block_dim == sum(aval.shape[dimension] for aval in ctx.avals_in):
    # Handle special case if the block contains all of the concatenated
    # array.
    return jax.lax.concatenate(args, dimension=dimension)

  num_blocks = []
  for aval in ctx.avals_in:
    assert isinstance(aval, core.ShapedArray)
    if aval.shape[dimension] % block_dim != 0:
      raise ValueError(
          f'Shape along concat dimension {dimension} must be divisible by the'
          f' block shape {block_shape[dimension]} for all children. Got shape'

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove Element entries from the output BlockSpec; use int block sizes
Compute the concatenation outside the kernel and write blocks individually
Use jnp.stack-free manual assignment of slices into the output block

Example fix

# before
out_ref[...] = jnp.concatenate([a, b], axis=0)  # out uses Element blocks

# after
out_ref[...] = jnp.concatenate([a, b], axis=0)  # BlockSpec uses int blocks
# BlockSpec(block_shape=(BS, ...)) instead of (Element(...), ...)

Defensive patterns

Strategy: validation

Validate before calling

import jax._src.pallas.pallas_core as pc
assert not any(isinstance(b, pc.Element) for b in block_spec.block_shape), 'concat output cannot use Element blocks'

Type guard

def concat_safe_block_spec(bs):
    import jax._src.pallas.pallas_core as pc
    return not any(isinstance(b, pc.Element) for b in bs.block_shape)

Prevention

Avoid Element blocks on outputs receiving concatenated values
Write blocks individually instead of concatenating into Element-blocked refs

When it happens

Trigger: Using jnp.concatenate / lax.concatenate on values that flow into a pallas output ref blocked with pallas_core.Element along any dimension.

Common situations: Building outputs by concatenating tiles inside kernels with fine-grained Element block specs; upgrading JAX where Element blocks became available but concat support lagged.

Related errors

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