jax-ml/jax · error · NotImplementedError

Only power-of-2 num parts supported.

Error message

Only power-of-2 num parts supported.

What it means

Split is lowered by recursively halving with tt.split, which only divides into 2. Therefore the number of parts must be a power of 2; len(sizes) that is not a power of two (e.g. 3) raises NotImplementedError.

Source

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

@register_lowering(jax._src.lax.lax.unstack_p)
def _unstack_lowering_rule(ctx: LoweringRuleContext, x, *, axis):
  [x_aval] = ctx.avals_in
  if x_aval.shape[axis] != 2:
    raise NotImplementedError("Only unstack of size 2 is supported in Triton.")
  if axis != x_aval.ndim - 1:
    raise NotImplementedError("Only unstack along the last dimension is supported in Triton.")

  x = _ensure_ir_value(x, x_aval)
  return tuple(tt_dialect.split(x))


@register_lowering(lax.split_p)
def _split_lowering_rule(ctx: LoweringRuleContext, x, *, sizes, axis):
  pass
  # TODO(cjfj): Add support for larger powers of 2.
  num_parts = len(sizes)
  if num_parts != pallas_utils.next_power_of_2(num_parts):
    raise NotImplementedError("Only power-of-2 num parts supported.")
  if any(size != sizes[0] for size in sizes):
    raise NotImplementedError("Only equal-sized splits are supported.")

  def split_into_2(x):
    shape = ir.RankedTensorType(x.type).shape
    x = _reshape(x, shape[:axis] + [2, shape[axis] // 2] + shape[axis + 1 :])
    permutation = tuple(d for d in range(len(shape) + 1) if d != axis) + (axis,)
    return tuple(tt_dialect.split(tt_dialect.trans(x, permutation)))

  x_parts: tuple[ir.Value, ...] = (x,)
  while len(x_parts) < num_parts:
    x_parts = sum(map(split_into_2, x_parts), ())
  return x_parts


def _compute_offsets_from_indices(
    block_info: BlockInfo, nd_indexer: NDIndexer
) -> ir.Value:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad or regroup so the number of sections is 1, 2, 4, 8... and split accordingly
2. Use slicing with lax.slice_in_dim per section instead of lax.split
3. Move the split out of the kernel

Example fix

// before
parts = jax.lax.split(x, [2, 2, 2], axis=0)

// after
parts = (x[0:2], x[2:4], x[4:6])  # or lax.slice_in_dim per part

def is_pow2(n: int) -> bool:
    return n > 0 and (n & (n - 1)) == 0
assert is_pow2(len(sizes)), 'number of split sections must be a power of 2 in-kernel'

def pow2_sections(sizes) -> bool:
    n = len(sizes)
    return n > 0 and (n & (n - 1)) == 0

try:
    parts = jax.lax.split(x, sizes, axis=ax)
except NotImplementedError:
    parts = tuple(lax.slice_in_dim(x, s, e, axis=ax) for s, e in zip(starts, ends))

When it happens

Trigger: jax.lax.split(x, sizes=[2, 2, 2]) or any split producing a non-power-of-2 number of sections inside a Triton Pallas kernel.

Common situations: Splitting blocks into 3 chunks for pipelining; ported code from XLA where arbitrary split sizes are supported.

Related errors

• Only equal-sized splits are supported.
• cannot cast {src} to {dst_type}
• Only 2-argument concatenate is supported.
• Only concatenate along the last dimension is supported.
• Only arguments with shape [..., 1] are supported.

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