{"record":{"id":"f33a535aeb0d7293","repo":"jax-ml/jax","slug":"only-equal-sized-splits-are-supported","errorCode":null,"errorMessage":"Only equal-sized splits are supported.","messagePattern":"Only equal-sized splits are supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1874,"sourceCode":"  [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:\n  full_shape = block_info.full_shape_dtype.shape\n  num_squeezed_dims = sum(isinstance(b, pallas_core.Squeezed)","sourceCodeStart":1856,"sourceCodeEnd":1892,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1856-L1892","documentation":"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.","triggerScenarios":"jax.lax.split(x, sizes=[4, 2]) or any sizes list with differing values inside a Triton Pallas kernel.","commonSituations":"Splitting ragged workloads or boundary chunks (last block smaller than the rest); variable sequence lengths in attention-style kernels.","solutions":["Pad the tensor so all sections are equal, split, then trim the results","Use lax.slice_in_dim / indexing to extract unequal sections directly","Redesign block sizes so sections are uniform"],"exampleFix":"// before\nparts = jax.lax.split(x, [4, 2], axis=0)\n\n// after\nparts = (lax.slice_in_dim(x, 0, 4, axis=0), lax.slice_in_dim(x, 4, 6, axis=0))","handlingStrategy":"validation","validationCode":"assert len(set(sizes)) == 1, 'in-kernel split requires equal section sizes'","typeGuard":"def equal_sections(sizes) -> bool:\n    return len(set(sizes)) == 1","tryCatchPattern":"try:\n    parts = jax.lax.split(x, sizes, axis=ax)\nexcept NotImplementedError:\n    cuts = [0] + list(itertools.accumulate(sizes))\n    parts = tuple(lax.slice_in_dim(x, cuts[i], cuts[i+1], axis=ax) for i in range(len(sizes)))","preventionTips":["Pad to uniform section sizes before splitting","Use slicing for ragged splits","Keep boundary chunks out of the kernel"],"tags":["jax","triton","pallas","split","shape","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"}