{"record":{"id":"b8f3e874a810f516","repo":"jax-ml/jax","slug":"only-power-of-2-num-parts-supported","errorCode":null,"errorMessage":"Only power-of-2 num parts supported.","messagePattern":"Only power-of-2 num parts supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1872,"sourceCode":"@register_lowering(jax._src.lax.lax.unstack_p)\ndef _unstack_lowering_rule(ctx: LoweringRuleContext, x, *, axis):\n  [x_aval] = ctx.avals_in\n  if x_aval.shape[axis] != 2:\n    raise NotImplementedError(\"Only unstack of size 2 is supported in Triton.\")\n  if axis != x_aval.ndim - 1:\n    raise NotImplementedError(\"Only unstack along the last dimension is supported in Triton.\")\n\n  x = _ensure_ir_value(x, x_aval)\n  return tuple(tt_dialect.split(x))\n\n\n@register_lowering(lax.split_p)\ndef _split_lowering_rule(ctx: LoweringRuleContext, x, *, sizes, axis):\n  pass\n  # TODO(cjfj): Add support for larger powers of 2.\n  num_parts = len(sizes)\n  if num_parts != pallas_utils.next_power_of_2(num_parts):\n    raise NotImplementedError(\"Only power-of-2 num parts supported.\")\n  if any(size != sizes[0] for size in sizes):\n    raise NotImplementedError(\"Only equal-sized splits are supported.\")\n\n  def split_into_2(x):\n    shape = ir.RankedTensorType(x.type).shape\n    x = _reshape(x, shape[:axis] + [2, shape[axis] // 2] + shape[axis + 1 :])\n    permutation = tuple(d for d in range(len(shape) + 1) if d != axis) + (axis,)\n    return tuple(tt_dialect.split(tt_dialect.trans(x, permutation)))\n\n  x_parts: tuple[ir.Value, ...] = (x,)\n  while len(x_parts) < num_parts:\n    x_parts = sum(map(split_into_2, x_parts), ())\n  return x_parts\n\n\ndef _compute_offsets_from_indices(\n    block_info: BlockInfo, nd_indexer: NDIndexer\n) -> ir.Value:","sourceCodeStart":1854,"sourceCodeEnd":1890,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1854-L1890","documentation":"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.","triggerScenarios":"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.","commonSituations":"Splitting blocks into 3 chunks for pipelining; ported code from XLA where arbitrary split sizes are supported.","solutions":["Pad or regroup so the number of sections is 1, 2, 4, 8... and split accordingly","Use slicing with lax.slice_in_dim per section instead of lax.split","Move the split out of the kernel"],"exampleFix":"// before\nparts = jax.lax.split(x, [2, 2, 2], axis=0)\n\n// after\nparts = (x[0:2], x[2:4], x[4:6])  # or lax.slice_in_dim per part","handlingStrategy":"validation","validationCode":"def is_pow2(n: int) -> bool:\n    return n > 0 and (n & (n - 1)) == 0\nassert is_pow2(len(sizes)), 'number of split sections must be a power of 2 in-kernel'","typeGuard":"def pow2_sections(sizes) -> bool:\n    n = len(sizes)\n    return n > 0 and (n & (n - 1)) == 0","tryCatchPattern":"try:\n    parts = jax.lax.split(x, sizes, axis=ax)\nexcept NotImplementedError:\n    parts = tuple(lax.slice_in_dim(x, s, e, axis=ax) for s, e in zip(starts, ends))","preventionTips":["Use power-of-2 section counts for in-kernel splits","Fall back to slice_in_dim for arbitrary sectioning","Design grid/block sizes as powers of 2"],"tags":["jax","triton","pallas","split","not-implemented"],"backgroundTag":"unsupported-operation-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}