jax-ml/jax · error · NotImplementedError

Only equal-sized splits are supported.

Error message

Only equal-sized splits are supported.

What it means

The Triton split lowering works by recursively applying tt.split, so the number of parts must be a power of two (and sizes equal). A split into a non-power-of-two number of parts fails the next_power_of_2 check and raises.

Solutions

1. Pad the tensor so all sections are equal, split, then trim the results
2. Use lax.slice_in_dim / indexing to extract unequal sections directly
3. Redesign block sizes so sections are uniform

Example fix

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

// after
parts = (lax.slice_in_dim(x, 0, 4, axis=0), lax.slice_in_dim(x, 4, 6, axis=0))

assert len(set(sizes)) == 1, 'in-kernel split requires equal section sizes'

def equal_sections(sizes) -> bool:
    return len(set(sizes)) == 1

try:
    parts = jax.lax.split(x, sizes, axis=ax)
except NotImplementedError:
    cuts = [0] + list(itertools.accumulate(sizes))
    parts = tuple(lax.slice_in_dim(x, cuts[i], cuts[i+1], axis=ax) for i in range(len(sizes)))

When it happens

Trigger: jax.lax.split(x, sizes=[4, 2]) or any sizes list with differing values inside a Triton Pallas kernel.

Common situations: Splitting ragged workloads or boundary chunks (last block smaller than the rest); variable sequence lengths in attention-style kernels.

Related errors

• Only arguments with shape [..., 1] are supported.
• Only power-of-2 num parts supported.
• a and b must be 2D, but got
• cannot cast to
• cmp must be scalar.

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

Appendix: source

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

  [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:
  full_shape = block_info.full_shape_dtype.shape
  num_squeezed_dims = sum(isinstance(b, pallas_core.Squeezed)

View on GitHub (pinned to 1e1c6a8fc0)