jax-ml/jax · error · NotImplementedError

Only equal-sized splits are supported.

Error message

Only equal-sized splits are supported.

What it means

The recursive tt.split lowering requires every section to have equal length, since each halving step divides the axis evenly. Any split with unequal sizes raises NotImplementedError.

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)

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.
• cannot cast {src} to {dst_type}
• Only 2-argument concatenate is supported.
• Only concatenate along the last dimension is supported.

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